internal/storage/sqlite/cooldown_auth_error_test.go

package sqlite_test

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"ccLoad/internal/model"
	"ccLoad/internal/storage"
	"ccLoad/internal/util"
)

// TestAuthErrorInitialCooldown 验证401/403错误的初始冷却时间为5分钟
func TestAuthErrorInitialCooldown(t *testing.T) {
	tests := []struct {
		name           string
		statusCode     int
		expectedMinDur time.Duration
		expectedMaxDur time.Duration
	}{
		{
			name:           "401未认证错误-初始冷却5分钟",
			statusCode:     401,
			expectedMinDur: 5 * time.Minute,
			expectedMaxDur: 5 * time.Minute,
		},
		{
			name:           "403禁止访问错误-初始冷却5分钟",
			statusCode:     403,
			expectedMinDur: 5 * time.Minute,
			expectedMaxDur: 5 * time.Minute,
		},
		{
			name:           "429限流错误-初始冷却1分钟",
			statusCode:     429,
			expectedMinDur: time.Minute,
			expectedMaxDur: time.Minute,
		},
		{
			name:           "500服务器错误-初始冷却2分钟",
			statusCode:     500,
			expectedMinDur: 2 * time.Minute,
			expectedMaxDur: 2 * time.Minute,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// 创建临时测试数据库
			store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
			defer cleanup()

			ctx := context.Background()
			now := time.Now()

			// 创建测试渠道
			cfg := &model.Config{
				Name:    "test-channel",
				URL:     "https://api.example.com",
				Enabled: true,
			}
			created, err := store.CreateConfig(ctx, cfg)
			if err != nil {
				t.Fatalf("创建测试渠道失败: %v", err)
			}

			// 触发首次错误冷却
			duration, err := store.BumpChannelCooldown(ctx, created.ID, now, tt.statusCode)
			if err != nil {
				t.Fatalf("BumpCooldownOnError失败: %v", err)
			}

			// 验证冷却时长
			if duration < tt.expectedMinDur || duration > tt.expectedMaxDur {
				t.Errorf("状态码%d的初始冷却时间错误: 期望%v，实际%v",
					tt.statusCode, tt.expectedMinDur, duration)
			}

			// 验证数据库中的冷却截止时间
			until, exists := getChannelCooldownUntil(ctx, store, created.ID)
			if !exists {
				t.Fatal("冷却记录不存在")
			}

			actualDuration := until.Sub(now)
			tolerance := 1 * time.Second // 允许1秒误差（考虑测试执行耗时）

			if actualDuration < tt.expectedMinDur-tolerance || actualDuration > tt.expectedMaxDur+tolerance {
				t.Errorf("数据库冷却时间错误: 期望%v，实际%v",
					tt.expectedMinDur, actualDuration)
			}

			t.Logf("[INFO] 状态码%d: 初始冷却时间=%v（期望%v）",
				tt.statusCode, duration, tt.expectedMinDur)
		})
	}
}

// TestAuthErrorExponentialBackoff 验证401/403错误的指数退避机制
func TestAuthErrorExponentialBackoff(t *testing.T) {
	store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
	defer cleanup()

	ctx := context.Background()
	now := time.Now()

	// 创建测试渠道
	cfg := &model.Config{
		Name:    "test-channel-backoff",
		URL:     "https://api.example.com",
		Enabled: true,
	}
	created, err := store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 预期的退避序列：5min -> 10min -> 20min -> 30min (上限)
	expectedSequence := []time.Duration{
		5 * time.Minute,  // 首次401错误
		10 * time.Minute, // 第二次错误（5min * 2）
		20 * time.Minute, // 第三次错误（10min * 2）
		30 * time.Minute, // 第四次错误（20min * 2，但达到上限）
		30 * time.Minute, // 第五次错误（保持上限）
	}

	for i, expected := range expectedSequence {
		// 触发401错误
		duration, err := store.BumpChannelCooldown(ctx, created.ID, now, 401)
		if err != nil {
			t.Fatalf("第%d次BumpCooldownOnError失败: %v", i+1, err)
		}

		// 验证冷却时长
		tolerance := 100 * time.Millisecond
		if duration < expected-tolerance || duration > expected+tolerance {
			t.Errorf("第%d次错误的冷却时间错误: 期望%v，实际%v",
				i+1, expected, duration)
		}

		t.Logf("[INFO] 第%d次401错误: 冷却时间=%v（期望%v）",
			i+1, duration, expected)

		// 更新now模拟时间推移（否则会被当作同一次错误）
		now = now.Add(expected + 1*time.Second)
	}
}

// TestKeyLevelAuthErrorCooldown 验证Key级别的401/403错误冷却
func TestKeyLevelAuthErrorCooldown(t *testing.T) {
	store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
	defer cleanup()

	ctx := context.Background()
	now := time.Now()

	// 创建多Key渠道
	cfg := &model.Config{
		Name:    "multi-key-channel",
		URL:     "https://api.example.com",
		Enabled: true,
	}
	created, err := store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 创建3个API Keys
	keyNames := []string{"sk-key1", "sk-key2", "sk-key3"}
	keys := make([]*model.APIKey, len(keyNames))
	for i, key := range keyNames {
		keys[i] = &model.APIKey{
			ChannelID:   created.ID,
			KeyIndex:    i,
			APIKey:      key,
			KeyStrategy: model.KeyStrategySequential,
		}
	}
	if err = store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("批量创建API Keys失败: %v", err)
	}

	// 测试Key 0的401错误冷却
	duration, err := store.BumpKeyCooldown(ctx, created.ID, 0, now, 401)
	if err != nil {
		t.Fatalf("BumpKeyCooldownOnError失败: %v", err)
	}

	// 验证初始冷却时间为5分钟
	expectedDuration := 5 * time.Minute
	tolerance := 1 * time.Second // 允许1秒误差（考虑测试执行耗时）
	if duration < expectedDuration-tolerance || duration > expectedDuration+tolerance {
		t.Errorf("Key级401错误初始冷却时间错误: 期望%v，实际%v",
			expectedDuration, duration)
	}

	// 验证数据库中的Key冷却记录
	until, exists := getKeyCooldownUntil(ctx, store, created.ID, 0)
	if !exists {
		t.Fatal("Key冷却记录不存在")
	}

	actualDuration := until.Sub(now)
	if actualDuration < expectedDuration-tolerance || actualDuration > expectedDuration+tolerance {
		t.Errorf("数据库Key冷却时间错误: 期望%v，实际%v",
			expectedDuration, actualDuration)
	}

	t.Logf("[INFO] Key级401错误: 初始冷却时间=%v（期望%v）",
		duration, expectedDuration)
}

// TestMixedErrorCodesCooldown 验证不同错误码混合场景的冷却行为
func TestMixedErrorCodesCooldown(t *testing.T) {
	store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
	defer cleanup()

	ctx := context.Background()
	now := time.Now()

	// 创建测试渠道
	cfg := &model.Config{
		Name:    "mixed-errors-channel",
		URL:     "https://api.example.com",
		Enabled: true,
	}
	created, err := store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 场景：先遇到500错误（2分钟起），然后遇到401错误（应该还是5分钟）
	duration1, err := store.BumpChannelCooldown(ctx, created.ID, now, 500)
	if err != nil {
		t.Fatalf("首次500错误失败: %v", err)
	}

	if duration1 != 2*time.Minute {
		t.Errorf("500错误初始冷却时间错误: 期望2分钟，实际%v", duration1)
	}

	// 模拟时间推移后遇到401错误
	now2 := now.Add(3 * time.Minute)
	duration2, err := store.BumpChannelCooldown(ctx, created.ID, now2, 401)
	if err != nil {
		t.Fatalf("后续401错误失败: %v", err)
	}

	// 因为之前有2分钟的冷却记录，新的401错误应该基于历史记录进行指数退避
	// 预期: 2min * 2 = 4min（但401首次应该是5分钟）
	// 实际逻辑：有历史记录则基于历史翻倍，无历史则按状态码初始化
	// 这里因为有历史duration_ms，所以是翻倍逻辑：2min * 2 = 4min
	expectedDuration := 4 * time.Minute
	tolerance := 100 * time.Millisecond

	if duration2 < expectedDuration-tolerance || duration2 > expectedDuration+tolerance {
		t.Errorf("混合错误场景冷却时间错误: 期望%v，实际%v",
			expectedDuration, duration2)
	}

	t.Logf("[INFO] 500错误(2min) → 401错误(%v) - 使用指数退避而非重置", duration2)
}

// TestConcurrentCooldownUpdates 验证并发场景下冷却机制的数据一致性
func TestConcurrentCooldownUpdates(t *testing.T) {
	if testing.Short() {
		t.Skip("跳过并发测试（使用 -short 标志）")
	}

	store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
	defer cleanup()

	ctx := context.Background()

	// 创建测试渠道
	cfg := &model.Config{
		Name:    "concurrent-test",
		URL:     "https://api.example.com",
		Enabled: true,
	}
	created, err := store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 并发触发10次401错误（足以验证并发安全性）
	const concurrency = 10
	var wg sync.WaitGroup
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, _ = store.BumpChannelCooldown(ctx, created.ID, time.Now(), 401)
		}()
	}
	wg.Wait()

	// 验证数据一致性
	until, exists := getChannelCooldownUntil(ctx, store, created.ID)
	if !exists {
		t.Fatal("冷却记录不存在")
	}

	duration := time.Until(until)
	minDuration := util.AuthErrorInitialCooldown - 1*time.Second
	maxDuration := util.MaxCooldownDuration + 1*time.Second

	if duration < minDuration || duration > maxDuration {
		t.Errorf("并发场景冷却时间异常: %v (期望范围: %v - %v)",
			duration, minDuration, maxDuration)
	}

	t.Logf("[INFO] 并发测试通过: %d个并发更新，最终冷却时间=%v", concurrency, duration)
}

// TestConcurrentKeyCooldownUpdates 验证Key级别并发冷却的数据一致性
func TestConcurrentKeyCooldownUpdates(t *testing.T) {
	if testing.Short() {
		t.Skip("跳过并发测试（使用 -short 标志）")
	}

	store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
	defer cleanup()

	ctx := context.Background()

	// 创建多Key渠道
	cfg := &model.Config{
		Name:    "concurrent-key-test",
		URL:     "https://api.example.com",
		Enabled: true,
	}
	created, err := store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 创建3个API Keys
	keyNames := []string{"sk-key1", "sk-key2", "sk-key3"}
	keys := make([]*model.APIKey, len(keyNames))
	for i, key := range keyNames {
		keys[i] = &model.APIKey{
			ChannelID:   created.ID,
			KeyIndex:    i,
			APIKey:      key,
			KeyStrategy: model.KeyStrategySequential,
		}
	}
	if err = store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("批量创建API Keys失败: %v", err)
	}

	// 使用信号量控制并发度为2，避免过多BUSY错误
	sem := make(chan struct{}, 2)
	var wg sync.WaitGroup
	var successCount int32

	// 每个Key更新3次，共9次操作
	for keyIndex := 0; keyIndex < 3; keyIndex++ {
		for i := 0; i < 3; i++ {
			wg.Add(1)
			go func(idx int) {
				defer wg.Done()
				sem <- struct{}{}
				defer func() { <-sem }()

				_, err := store.BumpKeyCooldown(ctx, created.ID, idx, time.Now(), 401)
				if err == nil {
					atomic.AddInt32(&successCount, 1)
				}
			}(keyIndex)
		}
	}
	wg.Wait()

	t.Logf("[INFO] 并发更新完成: 成功次数=%d/9", successCount)

	// 验证每个Key的冷却状态
	for keyIndex := 0; keyIndex < 3; keyIndex++ {
		until, exists := getKeyCooldownUntil(ctx, store, created.ID, keyIndex)
		if !exists {
			t.Errorf("Key %d 冷却记录不存在", keyIndex)
			continue
		}

		duration := time.Until(until)
		minDuration := util.AuthErrorInitialCooldown - 1*time.Second
		maxDuration := util.MaxCooldownDuration + 1*time.Second

		if duration < minDuration || duration > maxDuration {
			t.Errorf("Key %d 并发场景冷却时间异常: %v (期望范围: %v - %v)",
				keyIndex, duration, minDuration, maxDuration)
		}
	}
}

// TestRaceConditionDetection 竞态条件检测测试
// 使用 go test -race 运行此测试
func TestRaceConditionDetection(t *testing.T) {
	if testing.Short() {
		t.Skip("跳过竞态检测测试（使用 -short 标志）")
	}

	store, cleanup := setupSQLiteTestStore(t, "test-auth-error.db")
	defer cleanup()

	ctx := context.Background()

	cfg := &model.Config{
		Name:    "race-test",
		URL:     "https://api.example.com",
		Enabled: true,
	}
	created, err := store.CreateConfig(ctx, cfg)
	if err != nil {
		t.Fatalf("创建测试渠道失败: %v", err)
	}

	// 创建2个API Keys
	_ = store.CreateAPIKeysBatch(ctx, []*model.APIKey{
		{ChannelID: created.ID, KeyIndex: 0, APIKey: "sk-key1", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: created.ID, KeyIndex: 1, APIKey: "sk-key2", KeyStrategy: model.KeyStrategySequential},
	})

	// 并发场景：同时读写冷却状态（降低并发度）
	var wg sync.WaitGroup
	for i := 0; i < 5; i++ {
		wg.Add(3)

		// 写操作：更新渠道冷却
		go func() {
			defer wg.Done()
			_, _ = store.BumpChannelCooldown(ctx, created.ID, time.Now(), 401)
		}()

		// 写操作：更新Key冷却
		go func() {
			defer wg.Done()
			_, _ = store.BumpKeyCooldown(ctx, created.ID, 0, time.Now(), 401)
		}()

		// 读操作：获取渠道配置
		go func() {
			defer wg.Done()
			_, _ = store.GetConfig(ctx, created.ID)
		}()
	}

	wg.Wait()
	t.Log("[INFO] 竞态检测测试通过（使用 go test -race 运行以检测竞态条件）")
}

// setupAuthErrorTestStore 创建临时测试数据库（专用于认证错误测试）
// setupSQLiteTestStore 见 test_store_helpers_test.go
// getChannelCooldownUntil 获取渠道冷却截止时间（测试辅助函数）
func getChannelCooldownUntil(ctx context.Context, store storage.Store, channelID int64) (time.Time, bool) {
	cfg, err := store.GetConfig(ctx, channelID)
	if err != nil || cfg == nil {
		return time.Time{}, false
	}
	if cfg.CooldownUntil == 0 {
		return time.Time{}, false
	}
	until := time.Unix(cfg.CooldownUntil, 0)
	// 只有未过期的冷却才返回true
	return until, time.Now().Before(until)
}

// getKeyCooldownUntil 获取Key冷却截止时间（测试辅助函数）
func getKeyCooldownUntil(ctx context.Context, store storage.Store, channelID int64, keyIndex int) (time.Time, bool) {
	key, err := store.GetAPIKey(ctx, channelID, keyIndex)
	if err != nil || key == nil {
		return time.Time{}, false
	}
	if key.CooldownUntil == 0 {
		return time.Time{}, false
	}
	until := time.Unix(key.CooldownUntil, 0)
	// 只有未过期的冷却才返回true
	return until, time.Now().Before(until)
}