internal/storage/sql/cooldown_extras_test.go

package sql_test

import (
	"context"
	"testing"
	"time"

	"ccLoad/internal/model"
	sqlstore "ccLoad/internal/storage/sql"
)

func TestCooldown_GetKeyCooldownUntil_AndClearAll(t *testing.T) {
	store := newTestStore(t, "cooldown_extras.db")
	ctx := context.Background()

	ss := store.(*sqlstore.SQLStore)

	cfg, err := store.CreateConfig(ctx, &model.Config{
		Name:         "c1",
		URL:          "https://example.com",
		Priority:     1,
		Enabled:      true,
		ChannelType:  "openai",
		ModelEntries: []model.ModelEntry{{Model: "gpt-4o"}},
	})
	if err != nil {
		t.Fatalf("CreateConfig failed: %v", err)
	}

	if err := store.CreateAPIKeysBatch(ctx, []*model.APIKey{
		{ChannelID: cfg.ID, KeyIndex: 0, APIKey: "k0", KeyStrategy: model.KeyStrategySequential},
	}); err != nil {
		t.Fatalf("CreateAPIKeysBatch failed: %v", err)
	}

	until := time.Now().Add(2 * time.Minute)
	if err := store.SetKeyCooldown(ctx, cfg.ID, 0, until); err != nil {
		t.Fatalf("SetKeyCooldown failed: %v", err)
	}

	gotUntil, ok := ss.GetKeyCooldownUntil(ctx, cfg.ID, 0)
	if !ok || gotUntil.IsZero() {
		t.Fatalf("GetKeyCooldownUntil ok=%v until=%v, want true,non-zero", ok, gotUntil)
	}

	if err := ss.ClearAllKeyCooldowns(ctx, cfg.ID); err != nil {
		t.Fatalf("ClearAllKeyCooldowns failed: %v", err)
	}
	_, ok = ss.GetKeyCooldownUntil(ctx, cfg.ID, 0)
	if ok {
		t.Fatalf("expected key cooldown cleared")
	}
}