package kiro import ( "sync" "testing" "time" ) func TestNewCooldownManager(t *testing.T) { cm := NewCooldownManager() if cm == nil { t.Fatal("expected non-nil CooldownManager") } if cm.cooldowns == nil { t.Error("expected non-nil cooldowns map") } if cm.reasons == nil { t.Error("expected non-nil reasons map") } } func TestSetCooldown(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Minute, CooldownReason429) if !cm.IsInCooldown("token1") { t.Error("expected token to be in cooldown") } if cm.GetCooldownReason("token1") != CooldownReason429 { t.Errorf("expected reason %s, got %s", CooldownReason429, cm.GetCooldownReason("token1")) } } func TestIsInCooldown_NotSet(t *testing.T) { cm := NewCooldownManager() if cm.IsInCooldown("nonexistent") { t.Error("expected non-existent token to not be in cooldown") } } func TestIsInCooldown_Expired(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Millisecond, CooldownReason429) time.Sleep(10 * time.Millisecond) if cm.IsInCooldown("token1") { t.Error("expected expired cooldown to return false") } } func TestGetRemainingCooldown(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Second, CooldownReason429) remaining := cm.GetRemainingCooldown("token1") if remaining <= 0 || remaining > 1*time.Second { t.Errorf("expected remaining cooldown between 0 and 1s, got %v", remaining) } } func TestGetRemainingCooldown_NotSet(t *testing.T) { cm := NewCooldownManager() remaining := cm.GetRemainingCooldown("nonexistent") if remaining != 0 { t.Errorf("expected 0 remaining for non-existent, got %v", remaining) } } func TestGetRemainingCooldown_Expired(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Millisecond, CooldownReason429) time.Sleep(10 * time.Millisecond) remaining := cm.GetRemainingCooldown("token1") if remaining != 0 { t.Errorf("expected 0 remaining for expired, got %v", remaining) } } func TestGetCooldownReason(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Minute, CooldownReasonSuspended) reason := cm.GetCooldownReason("token1") if reason != CooldownReasonSuspended { t.Errorf("expected reason %s, got %s", CooldownReasonSuspended, reason) } } func TestGetCooldownReason_NotSet(t *testing.T) { cm := NewCooldownManager() reason := cm.GetCooldownReason("nonexistent") if reason != "" { t.Errorf("expected empty reason for non-existent, got %s", reason) } } func TestClearCooldown(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Minute, CooldownReason429) cm.ClearCooldown("token1") if cm.IsInCooldown("token1") { t.Error("expected cooldown to be cleared") } if cm.GetCooldownReason("token1") != "" { t.Error("expected reason to be cleared") } } func TestClearCooldown_NonExistent(t *testing.T) { cm := NewCooldownManager() cm.ClearCooldown("nonexistent") } func TestCleanupExpired(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("expired1", 1*time.Millisecond, CooldownReason429) cm.SetCooldown("expired2", 1*time.Millisecond, CooldownReason429) cm.SetCooldown("active", 1*time.Hour, CooldownReason429) time.Sleep(10 * time.Millisecond) cm.CleanupExpired() if cm.GetCooldownReason("expired1") != "" { t.Error("expected expired1 to be cleaned up") } if cm.GetCooldownReason("expired2") != "" { t.Error("expected expired2 to be cleaned up") } if cm.GetCooldownReason("active") != CooldownReason429 { t.Error("expected active to remain") } } func TestCalculateCooldownFor429_FirstRetry(t *testing.T) { duration := CalculateCooldownFor429(0) if duration != DefaultShortCooldown { t.Errorf("expected %v for retry 0, got %v", DefaultShortCooldown, duration) } } func TestCalculateCooldownFor429_Exponential(t *testing.T) { d1 := CalculateCooldownFor429(1) d2 := CalculateCooldownFor429(2) if d2 <= d1 { t.Errorf("expected d2 > d1, got d1=%v, d2=%v", d1, d2) } } func TestCalculateCooldownFor429_MaxCap(t *testing.T) { duration := CalculateCooldownFor429(10) if duration > MaxShortCooldown { t.Errorf("expected max %v, got %v", MaxShortCooldown, duration) } } func TestCalculateCooldownUntilNextDay(t *testing.T) { duration := CalculateCooldownUntilNextDay() if duration <= 0 || duration > 24*time.Hour { t.Errorf("expected duration between 0 and 24h, got %v", duration) } } func TestCooldownManager_ConcurrentAccess(t *testing.T) { cm := NewCooldownManager() const numGoroutines = 50 const numOperations = 100 var wg sync.WaitGroup wg.Add(numGoroutines) for i := 0; i < numGoroutines; i++ { go func(id int) { defer wg.Done() tokenKey := "token" + string(rune('a'+id%10)) for j := 0; j < numOperations; j++ { switch j % 6 { case 0: cm.SetCooldown(tokenKey, time.Duration(j)*time.Millisecond, CooldownReason429) case 1: cm.IsInCooldown(tokenKey) case 2: cm.GetRemainingCooldown(tokenKey) case 3: cm.GetCooldownReason(tokenKey) case 4: cm.ClearCooldown(tokenKey) case 5: cm.CleanupExpired() } } }(i) } wg.Wait() } func TestCooldownReasonConstants(t *testing.T) { if CooldownReason429 != "rate_limit_exceeded" { t.Errorf("unexpected CooldownReason429: %s", CooldownReason429) } if CooldownReasonSuspended != "account_suspended" { t.Errorf("unexpected CooldownReasonSuspended: %s", CooldownReasonSuspended) } if CooldownReasonQuotaExhausted != "quota_exhausted" { t.Errorf("unexpected CooldownReasonQuotaExhausted: %s", CooldownReasonQuotaExhausted) } } func TestDefaultConstants(t *testing.T) { if DefaultShortCooldown != 1*time.Minute { t.Errorf("unexpected DefaultShortCooldown: %v", DefaultShortCooldown) } if MaxShortCooldown != 5*time.Minute { t.Errorf("unexpected MaxShortCooldown: %v", MaxShortCooldown) } if LongCooldown != 24*time.Hour { t.Errorf("unexpected LongCooldown: %v", LongCooldown) } } func TestSetCooldown_OverwritesPrevious(t *testing.T) { cm := NewCooldownManager() cm.SetCooldown("token1", 1*time.Hour, CooldownReason429) cm.SetCooldown("token1", 1*time.Minute, CooldownReasonSuspended) reason := cm.GetCooldownReason("token1") if reason != CooldownReasonSuspended { t.Errorf("expected reason to be overwritten to %s, got %s", CooldownReasonSuspended, reason) } remaining := cm.GetRemainingCooldown("token1") if remaining > 1*time.Minute { t.Errorf("expected remaining <= 1 minute, got %v", remaining) } }