anthropic-proxy/internal/ratelimit/tracker_test.go

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 }