package ratelimit import ( "net/http" "testing" "time" ) func TestNewTracker(t *testing.T) { called := false tr := NewTracker(func() string { called = true return "tok" }) if tr == nil { t.Fatal("NewTracker returned nil") } // tokenFn stored but not called during construction if called { t.Error("tokenFn should not be called by NewTracker") } // Invoke to verify it's wired if got := tr.tokenFn(); got != "tok" { t.Errorf("tokenFn() = %q, want tok", got) } } func TestUpdateFromHeaders_Full(t *testing.T) { tr := NewTracker(func() string { return "" }) h := http.Header{} h.Set("Anthropic-Ratelimit-Unified-5h-Utilization", "0.42") h.Set("Anthropic-Ratelimit-Unified-5h-Reset", "1700000000") h.Set("Anthropic-Ratelimit-Unified-7d-Utilization", "0.75") h.Set("Anthropic-Ratelimit-Unified-7d-Reset", "1700100000") tr.UpdateFromHeaders(h) fh := tr.FiveHour() if fh.Utilization != 42.0 { t.Errorf("FiveHour.Utilization = %f, want 42.0", fh.Utilization) } wantReset5h := time.Unix(1700000000, 0).UTC().Truncate(time.Minute) if !fh.ResetsAt.Equal(wantReset5h) { t.Errorf("FiveHour.ResetsAt = %v, want %v", fh.ResetsAt, wantReset5h) } sd := tr.SevenDay() if sd.Utilization != 75.0 { t.Errorf("SevenDay.Utilization = %f, want 75.0", sd.Utilization) } wantReset7d := time.Unix(1700100000, 0).UTC().Truncate(time.Minute) if !sd.ResetsAt.Equal(wantReset7d) { t.Errorf("SevenDay.ResetsAt = %v, want %v", sd.ResetsAt, wantReset7d) } } func TestUpdateFromHeaders_Partial(t *testing.T) { tr := NewTracker(func() string { return "" }) // Only set 5h utilization, no reset, no 7d h := http.Header{} h.Set("Anthropic-Ratelimit-Unified-5h-Utilization", "0.33") tr.UpdateFromHeaders(h) fh := tr.FiveHour() if fh.Utilization != 33.0 { t.Errorf("FiveHour.Utilization = %f, want 33.0", fh.Utilization) } if !fh.ResetsAt.IsZero() { t.Errorf("FiveHour.ResetsAt should be zero, got %v", fh.ResetsAt) } sd := tr.SevenDay() if sd.Utilization != 0 { t.Errorf("SevenDay.Utilization = %f, want 0", sd.Utilization) } } func TestUpdateFromHeaders_Missing(t *testing.T) { tr := NewTracker(func() string { return "" }) // Pre-set some state tr.mu.Lock() tr.fiveHour.Utilization = 50.0 tr.mu.Unlock() // Update with empty headers — should not change state tr.UpdateFromHeaders(http.Header{}) fh := tr.FiveHour() if fh.Utilization != 50.0 { t.Errorf("FiveHour.Utilization = %f, want 50.0 (unchanged)", fh.Utilization) } } func TestUpdateFromHeaders_InvalidValues(t *testing.T) { tr := NewTracker(func() string { return "" }) h := http.Header{} h.Set("Anthropic-Ratelimit-Unified-5h-Utilization", "not-a-number") h.Set("Anthropic-Ratelimit-Unified-5h-Reset", "not-a-timestamp") tr.UpdateFromHeaders(h) fh := tr.FiveHour() if fh.Utilization != 0 { t.Errorf("Utilization should stay 0 for invalid input, got %f", fh.Utilization) } if !fh.ResetsAt.IsZero() { t.Errorf("ResetsAt should stay zero for invalid input, got %v", fh.ResetsAt) } } func TestSonnet_Snapshot(t *testing.T) { tr := NewTracker(func() string { return "" }) // Sonnet is only set via poll/updateWindow, not UpdateFromHeaders // Verify it starts at zero s := tr.Sonnet() if s.Utilization != 0 { t.Errorf("Sonnet.Utilization = %f, want 0", s.Utilization) } if !s.ResetsAt.IsZero() { t.Errorf("Sonnet.ResetsAt should be zero, got %v", s.ResetsAt) } } func TestExtra_Default(t *testing.T) { tr := NewTracker(func() string { return "" }) extra := tr.Extra() if extra.IsEnabled { t.Error("Extra.IsEnabled should be false by default") } if extra.MonthlyLimit != nil { t.Error("Extra.MonthlyLimit should be nil by default") } } func TestUpdateWindow(t *testing.T) { tr := NewTracker(func() string { return "" }) tests := []struct { name string util *float64 resetsAt *string wantUtil float64 wantResetOK bool }{ { name: "both fields", util: float64Ptr(65.5), resetsAt: stringPtr("2024-01-15T10:30:45Z"), wantUtil: 65.5, wantResetOK: true, }, { name: "utilization only", util: float64Ptr(30.0), resetsAt: nil, wantUtil: 30.0, wantResetOK: false, }, { name: "reset only (RFC3339Nano)", util: nil, resetsAt: stringPtr("2024-06-01T12:00:00.123456789Z"), wantUtil: 0, wantResetOK: true, }, { name: "nil both", util: nil, resetsAt: nil, wantUtil: 0, wantResetOK: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { w := &Window{} rl := &RateLimit{ Utilization: tt.util, ResetsAt: tt.resetsAt, } tr.updateWindow(w, rl) if w.Utilization != tt.wantUtil { t.Errorf("Utilization = %f, want %f", w.Utilization, tt.wantUtil) } if tt.wantResetOK { if w.ResetsAt.IsZero() { t.Error("ResetsAt should be set") } // Verify truncation to minute if w.ResetsAt.Second() != 0 || w.ResetsAt.Nanosecond() != 0 { t.Errorf("ResetsAt not truncated to minute: %v", w.ResetsAt) } if w.ResetsAt.Location() != time.UTC { t.Errorf("ResetsAt not in UTC: %v", w.ResetsAt.Location()) } } else if tt.resetsAt == nil { if !w.ResetsAt.IsZero() { t.Errorf("ResetsAt should be zero when input is nil, got %v", w.ResetsAt) } } }) } } func TestUpdateWindow_InvalidTime(t *testing.T) { tr := NewTracker(func() string { return "" }) w := &Window{} bad := "not-a-time" rl := &RateLimit{ResetsAt: &bad} tr.updateWindow(w, rl) if !w.ResetsAt.IsZero() { t.Errorf("ResetsAt should stay zero for invalid time, got %v", w.ResetsAt) } } func TestPoll_SetsStateFromUsageResponse(t *testing.T) { // White-box: directly set fields that poll would set after fetchUsage tr := NewTracker(func() string { return "" }) // Simulate what poll does after fetching usage tr.mu.Lock() usage := &UsageResponse{ FiveHour: &RateLimit{Utilization: float64Ptr(55.5), ResetsAt: stringPtr("2024-03-01T08:00:00Z")}, SevenDay: &RateLimit{Utilization: float64Ptr(22.3), ResetsAt: stringPtr("2024-03-07T00:00:00Z")}, SevenDaySonnet: &RateLimit{Utilization: float64Ptr(10.0), ResetsAt: stringPtr("2024-03-07T00:00:00Z")}, ExtraUsage: &ExtraUsage{IsEnabled: true, MonthlyLimit: float64Ptr(100.0), UsedCredits: float64Ptr(42.5)}, } if usage.FiveHour != nil { tr.updateWindow(&tr.fiveHour, usage.FiveHour) } if usage.SevenDay != nil { tr.updateWindow(&tr.sevenDay, usage.SevenDay) } if usage.SevenDaySonnet != nil { tr.updateWindow(&tr.sonnet, usage.SevenDaySonnet) } if usage.ExtraUsage != nil { tr.extra = *usage.ExtraUsage } tr.mu.Unlock() fh := tr.FiveHour() if fh.Utilization != 55.5 { t.Errorf("FiveHour.Utilization = %f, want 55.5", fh.Utilization) } sd := tr.SevenDay() if sd.Utilization != 22.3 { t.Errorf("SevenDay.Utilization = %f, want 22.3", sd.Utilization) } sn := tr.Sonnet() if sn.Utilization != 10.0 { t.Errorf("Sonnet.Utilization = %f, want 10.0", sn.Utilization) } extra := tr.Extra() if !extra.IsEnabled { t.Error("Extra.IsEnabled = false, want true") } if extra.MonthlyLimit == nil || *extra.MonthlyLimit != 100.0 { t.Errorf("Extra.MonthlyLimit = %v, want 100.0", extra.MonthlyLimit) } if extra.UsedCredits == nil || *extra.UsedCredits != 42.5 { t.Errorf("Extra.UsedCredits = %v, want 42.5", extra.UsedCredits) } } func float64Ptr(f float64) *float64 { return &f } func stringPtr(s string) *string { return &s }