anthropic-proxy/internal/auth/selector_test.go

package auth

import (
	"testing"
	"time"
)

func TestNewPool(t *testing.T) {
	creds := []*Credential{
		{ID: "a", AccessToken: "tok-a"},
		{ID: "b", AccessToken: "tok-b"},
	}
	p := NewPool(creds)
	if p == nil {
		t.Fatal("NewPool returned nil")
	}
	if len(p.creds) != 2 {
		t.Errorf("pool has %d creds, want 2", len(p.creds))
	}
	if p.cursor != 0 {
		t.Errorf("initial cursor = %d, want 0", p.cursor)
	}
}

func TestPool_Pick_EmptyPool(t *testing.T) {
	p := NewPool(nil)
	_, err := p.Pick()
	if err == nil {
		t.Fatal("expected error from empty pool, got nil")
	}
	want := "no credentials available"
	if err.Error() != want {
		t.Errorf("error = %q, want %q", err.Error(), want)
	}
}

func TestPool_Pick_SingleCredential(t *testing.T) {
	cred := &Credential{ID: "only", AccessToken: "tok-only"}
	p := NewPool([]*Credential{cred})

	got, err := p.Pick()
	if err != nil {
		t.Fatalf("Pick() error = %v", err)
	}
	if got.ID != "only" {
		t.Errorf("Pick() returned cred ID %q, want %q", got.ID, "only")
	}

	// Picking again should return the same credential
	got2, err := p.Pick()
	if err != nil {
		t.Fatalf("second Pick() error = %v", err)
	}
	if got2.ID != "only" {
		t.Errorf("second Pick() returned cred ID %q, want %q", got2.ID, "only")
	}
}

func TestPool_Pick_RoundRobin(t *testing.T) {
	creds := []*Credential{
		{ID: "a"},
		{ID: "b"},
		{ID: "c"},
	}
	p := NewPool(creds)

	// Should cycle through a, b, c, a, b, c
	expected := []string{"a", "b", "c", "a", "b", "c"}
	for i, want := range expected {
		got, err := p.Pick()
		if err != nil {
			t.Fatalf("Pick() #%d error = %v", i, err)
		}
		if got.ID != want {
			t.Errorf("Pick() #%d = %q, want %q", i, got.ID, want)
		}
	}
}

func TestPool_Pick_SkipsCooldown(t *testing.T) {
	creds := []*Credential{
		{ID: "a"},
		{ID: "b", CooldownUntil: time.Now().Add(1 * time.Hour)},
		{ID: "c"},
	}
	p := NewPool(creds)

	// First pick: "a" (index 0, not on cooldown)
	got, err := p.Pick()
	if err != nil {
		t.Fatalf("Pick() #1 error = %v", err)
	}
	if got.ID != "a" {
		t.Errorf("Pick() #1 = %q, want %q", got.ID, "a")
	}

	// Second pick: cursor at 1, but "b" is on cooldown → skip to "c"
	got, err = p.Pick()
	if err != nil {
		t.Fatalf("Pick() #2 error = %v", err)
	}
	if got.ID != "c" {
		t.Errorf("Pick() #2 = %q, want %q", got.ID, "c")
	}

	// Third pick: cursor advanced past "c" to 0 → "a"
	got, err = p.Pick()
	if err != nil {
		t.Fatalf("Pick() #3 error = %v", err)
	}
	if got.ID != "a" {
		t.Errorf("Pick() #3 = %q, want %q", got.ID, "a")
	}
}

func TestPool_Pick_AllOnCooldown(t *testing.T) {
	future := time.Now().Add(1 * time.Hour)
	creds := []*Credential{
		{ID: "a", CooldownUntil: future},
		{ID: "b", CooldownUntil: future},
	}
	p := NewPool(creds)

	_, err := p.Pick()
	if err == nil {
		t.Fatal("expected error when all on cooldown, got nil")
	}
	want := "all 2 credentials are on cooldown"
	if err.Error() != want {
		t.Errorf("error = %q, want %q", err.Error(), want)
	}
}

func TestPool_MarkFailure(t *testing.T) {
	tests := []struct {
		name           string
		statusCode     int
		expectCooldown bool
		expectedDur    time.Duration
	}{
		{
			name:           "429 sets 30s cooldown",
			statusCode:     429,
			expectCooldown: true,
			expectedDur:    30 * time.Second,
		},
		{
			name:           "500 sets 5s cooldown",
			statusCode:     500,
			expectCooldown: true,
			expectedDur:    5 * time.Second,
		},
		{
			name:           "502 sets 5s cooldown",
			statusCode:     502,
			expectCooldown: true,
			expectedDur:    5 * time.Second,
		},
		{
			name:           "503 sets 5s cooldown",
			statusCode:     503,
			expectCooldown: true,
			expectedDur:    5 * time.Second,
		},
		{
			name:           "400 does NOT set cooldown",
			statusCode:     400,
			expectCooldown: false,
		},
		{
			name:           "401 does NOT set cooldown",
			statusCode:     401,
			expectCooldown: false,
		},
		{
			name:           "403 does NOT set cooldown",
			statusCode:     403,
			expectCooldown: false,
		},
		{
			name:           "404 does NOT set cooldown",
			statusCode:     404,
			expectCooldown: false,
		},
		{
			name:           "422 does NOT set cooldown",
			statusCode:     422,
			expectCooldown: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cred := &Credential{ID: "test"}
			p := NewPool([]*Credential{cred})

			before := time.Now()
			p.MarkFailure(cred, tt.statusCode)

			if tt.expectCooldown {
				if !cred.IsOnCooldown() {
					t.Errorf("expected cooldown after status %d", tt.statusCode)
				}
				// Verify approximate duration
				cred.mu.Lock()
				cooldownEnd := cred.CooldownUntil
				cred.mu.Unlock()

				lower := before.Add(tt.expectedDur)
				upper := time.Now().Add(tt.expectedDur)
				if cooldownEnd.Before(lower) || cooldownEnd.After(upper) {
					t.Errorf("CooldownUntil %v not in expected range [%v, %v]", cooldownEnd, lower, upper)
				}
			} else {
				if cred.IsOnCooldown() {
					t.Errorf("did not expect cooldown after status %d", tt.statusCode)
				}
			}
		})
	}
}

func TestPool_MarkSuccess(t *testing.T) {
	cred := &Credential{
		ID:            "test",
		CooldownUntil: time.Now().Add(1 * time.Hour),
	}
	p := NewPool([]*Credential{cred})

	if !cred.IsOnCooldown() {
		t.Fatal("precondition: expected credential to be on cooldown")
	}

	p.MarkSuccess(cred)

	if cred.IsOnCooldown() {
		t.Error("expected cooldown to be cleared after MarkSuccess")
	}
}

func TestPool_RoundRobinCursorAdvancement(t *testing.T) {
	creds := []*Credential{
		{ID: "0"},
		{ID: "1"},
		{ID: "2"},
	}
	p := NewPool(creds)

	// Verify cursor starts at 0
	if p.cursor != 0 {
		t.Fatalf("initial cursor = %d, want 0", p.cursor)
	}

	// Pick cred[0], cursor should advance to 1
	got, _ := p.Pick()
	if got.ID != "0" {
		t.Errorf("first pick = %q, want %q", got.ID, "0")
	}
	if p.cursor != 1 {
		t.Errorf("cursor after first pick = %d, want 1", p.cursor)
	}

	// Pick cred[1], cursor should advance to 2
	got, _ = p.Pick()
	if got.ID != "1" {
		t.Errorf("second pick = %q, want %q", got.ID, "1")
	}
	if p.cursor != 2 {
		t.Errorf("cursor after second pick = %d, want 2", p.cursor)
	}

	// Pick cred[2], cursor should wrap to 0
	got, _ = p.Pick()
	if got.ID != "2" {
		t.Errorf("third pick = %q, want %q", got.ID, "2")
	}
	if p.cursor != 0 {
		t.Errorf("cursor after third pick = %d, want 0 (wrap)", p.cursor)
	}
}

func TestPool_RoundRobinWithCooldownSkip(t *testing.T) {
	creds := []*Credential{
		{ID: "0"},
		{ID: "1", CooldownUntil: time.Now().Add(1 * time.Hour)},
		{ID: "2"},
	}
	p := NewPool(creds)

	// First pick: cred[0]
	got, _ := p.Pick()
	if got.ID != "0" {
		t.Errorf("first pick = %q, want %q", got.ID, "0")
	}
	// Cursor should be at 1
	if p.cursor != 1 {
		t.Errorf("cursor after first pick = %d, want 1", p.cursor)
	}

	// Second pick: cursor at 1, but cred[1] on cooldown → skip to cred[2]
	got, _ = p.Pick()
	if got.ID != "2" {
		t.Errorf("second pick = %q, want %q", got.ID, "2")
	}
	// Cursor should advance past cred[2] to 0
	if p.cursor != 0 {
		t.Errorf("cursor after second pick (skip) = %d, want 0", p.cursor)
	}

	// Third pick: cursor at 0, cred[0] available
	got, _ = p.Pick()
	if got.ID != "0" {
		t.Errorf("third pick = %q, want %q", got.ID, "0")
	}
	if p.cursor != 1 {
		t.Errorf("cursor after third pick = %d, want 1", p.cursor)
	}
}