package auth import ( "context" "errors" "sync" "testing" "time" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" ) func TestFillFirstSelectorPick_Deterministic(t *testing.T) { t.Parallel() selector := &FillFirstSelector{} auths := []*Auth{ {ID: "b"}, {ID: "a"}, {ID: "c"}, } got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if err != nil { t.Fatalf("Pick() error = %v", err) } if got == nil { t.Fatalf("Pick() auth = nil") } if got.ID != "a" { t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "a") } } func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) { t.Parallel() selector := &RoundRobinSelector{} auths := []*Auth{ {ID: "b"}, {ID: "a"}, {ID: "c"}, } want := []string{"a", "b", "c", "a", "b"} for i, id := range want { got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if err != nil { t.Fatalf("Pick() #%d error = %v", i, err) } if got == nil { t.Fatalf("Pick() #%d auth = nil", i) } if got.ID != id { t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id) } } } func TestRoundRobinSelectorPick_PriorityBuckets(t *testing.T) { t.Parallel() selector := &RoundRobinSelector{} auths := []*Auth{ {ID: "c", Attributes: map[string]string{"priority": "0"}}, {ID: "a", Attributes: map[string]string{"priority": "10"}}, {ID: "b", Attributes: map[string]string{"priority": "10"}}, } want := []string{"a", "b", "a", "b"} for i, id := range want { got, err := selector.Pick(context.Background(), "mixed", "", cliproxyexecutor.Options{}, auths) if err != nil { t.Fatalf("Pick() #%d error = %v", i, err) } if got == nil { t.Fatalf("Pick() #%d auth = nil", i) } if got.ID != id { t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id) } if got.ID == "c" { t.Fatalf("Pick() #%d unexpectedly selected lower priority auth", i) } } } func TestFillFirstSelectorPick_PriorityFallbackCooldown(t *testing.T) { t.Parallel() selector := &FillFirstSelector{} now := time.Now() model := "test-model" high := &Auth{ ID: "high", Attributes: map[string]string{"priority": "10"}, ModelStates: map[string]*ModelState{ model: { Status: StatusActive, Unavailable: true, NextRetryAfter: now.Add(30 * time.Minute), Quota: QuotaState{ Exceeded: true, }, }, }, } low := &Auth{ID: "low", Attributes: map[string]string{"priority": "0"}} got, err := selector.Pick(context.Background(), "mixed", model, cliproxyexecutor.Options{}, []*Auth{high, low}) if err != nil { t.Fatalf("Pick() error = %v", err) } if got == nil { t.Fatalf("Pick() auth = nil") } if got.ID != "low" { t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "low") } } func TestRoundRobinSelectorPick_Concurrent(t *testing.T) { selector := &RoundRobinSelector{} auths := []*Auth{ {ID: "b"}, {ID: "a"}, {ID: "c"}, } start := make(chan struct{}) var wg sync.WaitGroup errCh := make(chan error, 1) goroutines := 32 iterations := 100 for i := 0; i < goroutines; i++ { wg.Add(1) go func() { defer wg.Done() <-start for j := 0; j < iterations; j++ { got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if err != nil { select { case errCh <- err: default: } return } if got == nil { select { case errCh <- errors.New("Pick() returned nil auth"): default: } return } if got.ID == "" { select { case errCh <- errors.New("Pick() returned auth with empty ID"): default: } return } } }() } close(start) wg.Wait() select { case err := <-errCh: t.Fatalf("concurrent Pick() error = %v", err) default: } }