sdk/cliproxy/auth: update selector tests

sdk/cliproxy/auth/selector_test.go

@@ -2,7 +2,9 @@ package auth
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
+	"net/http"
 	"sync"
 	"testing"
 	"time"
@@ -175,3 +177,228 @@ func TestRoundRobinSelectorPick_Concurrent(t *testing.T) {
 	default:
 	}
 }
+
+func TestSelectorPick_AllCooldownReturnsModelCooldownError(t *testing.T) {
+	t.Parallel()
+
+	model := "test-model"
+	now := time.Now()
+	next := now.Add(60 * time.Second)
+	auths := []*Auth{
+		{
+			ID: "a",
+			ModelStates: map[string]*ModelState{
+				model: {
+					Status:         StatusActive,
+					Unavailable:    true,
+					NextRetryAfter: next,
+					Quota: QuotaState{
+						Exceeded:      true,
+						NextRecoverAt: next,
+					},
+				},
+			},
+		},
+		{
+			ID: "b",
+			ModelStates: map[string]*ModelState{
+				model: {
+					Status:         StatusActive,
+					Unavailable:    true,
+					NextRetryAfter: next,
+					Quota: QuotaState{
+						Exceeded:      true,
+						NextRecoverAt: next,
+					},
+				},
+			},
+		},
+	}
+
+	t.Run("mixed provider redacts provider field", func(t *testing.T) {
+		t.Parallel()
+
+		selector := &FillFirstSelector{}
+		_, err := selector.Pick(context.Background(), "mixed", model, cliproxyexecutor.Options{}, auths)
+		if err == nil {
+			t.Fatalf("Pick() error = nil")
+		}
+
+		var mce *modelCooldownError
+		if !errors.As(err, &mce) {
+			t.Fatalf("Pick() error = %T, want *modelCooldownError", err)
+		}
+		if mce.StatusCode() != http.StatusTooManyRequests {
+			t.Fatalf("StatusCode() = %d, want %d", mce.StatusCode(), http.StatusTooManyRequests)
+		}
+
+		headers := mce.Headers()
+		if got := headers.Get("Retry-After"); got == "" {
+			t.Fatalf("Headers().Get(Retry-After) = empty")
+		}
+
+		var payload map[string]any
+		if err := json.Unmarshal([]byte(mce.Error()), &payload); err != nil {
+			t.Fatalf("json.Unmarshal(Error()) error = %v", err)
+		}
+		rawErr, ok := payload["error"].(map[string]any)
+		if !ok {
+			t.Fatalf("Error() payload missing error object: %v", payload)
+		}
+		if got, _ := rawErr["code"].(string); got != "model_cooldown" {
+			t.Fatalf("Error().error.code = %q, want %q", got, "model_cooldown")
+		}
+		if _, ok := rawErr["provider"]; ok {
+			t.Fatalf("Error().error.provider exists for mixed provider: %v", rawErr["provider"])
+		}
+	})
+
+	t.Run("non-mixed provider includes provider field", func(t *testing.T) {
+		t.Parallel()
+
+		selector := &FillFirstSelector{}
+		_, err := selector.Pick(context.Background(), "gemini", model, cliproxyexecutor.Options{}, auths)
+		if err == nil {
+			t.Fatalf("Pick() error = nil")
+		}
+
+		var mce *modelCooldownError
+		if !errors.As(err, &mce) {
+			t.Fatalf("Pick() error = %T, want *modelCooldownError", err)
+		}
+
+		var payload map[string]any
+		if err := json.Unmarshal([]byte(mce.Error()), &payload); err != nil {
+			t.Fatalf("json.Unmarshal(Error()) error = %v", err)
+		}
+		rawErr, ok := payload["error"].(map[string]any)
+		if !ok {
+			t.Fatalf("Error() payload missing error object: %v", payload)
+		}
+		if got, _ := rawErr["provider"].(string); got != "gemini" {
+			t.Fatalf("Error().error.provider = %q, want %q", got, "gemini")
+		}
+	})
+}
+
+func TestIsAuthBlockedForModel_UnavailableWithoutNextRetryIsNotBlocked(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+	model := "test-model"
+	auth := &Auth{
+		ID: "a",
+		ModelStates: map[string]*ModelState{
+			model: {
+				Status:      StatusActive,
+				Unavailable: true,
+				Quota: QuotaState{
+					Exceeded: true,
+				},
+			},
+		},
+	}
+
+	blocked, reason, next := isAuthBlockedForModel(auth, model, now)
+	if blocked {
+		t.Fatalf("blocked = true, want false")
+	}
+	if reason != blockReasonNone {
+		t.Fatalf("reason = %v, want %v", reason, blockReasonNone)
+	}
+	if !next.IsZero() {
+		t.Fatalf("next = %v, want zero", next)
+	}
+}
+
+func TestFillFirstSelectorPick_ThinkingSuffixFallsBackToBaseModelState(t *testing.T) {
+	t.Parallel()
+
+	selector := &FillFirstSelector{}
+	now := time.Now()
+
+	baseModel := "test-model"
+	requestedModel := "test-model(high)"
+
+	high := &Auth{
+		ID:         "high",
+		Attributes: map[string]string{"priority": "10"},
+		ModelStates: map[string]*ModelState{
+			baseModel: {
+				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", requestedModel, 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_ThinkingSuffixSharesCursor(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{}
+	auths := []*Auth{
+		{ID: "b"},
+		{ID: "a"},
+	}
+
+	first, err := selector.Pick(context.Background(), "gemini", "test-model(high)", cliproxyexecutor.Options{}, auths)
+	if err != nil {
+		t.Fatalf("Pick() first error = %v", err)
+	}
+	second, err := selector.Pick(context.Background(), "gemini", "test-model(low)", cliproxyexecutor.Options{}, auths)
+	if err != nil {
+		t.Fatalf("Pick() second error = %v", err)
+	}
+	if first == nil || second == nil {
+		t.Fatalf("Pick() returned nil auth")
+	}
+	if first.ID != "a" {
+		t.Fatalf("Pick() first auth.ID = %q, want %q", first.ID, "a")
+	}
+	if second.ID != "b" {
+		t.Fatalf("Pick() second auth.ID = %q, want %q", second.ID, "b")
+	}
+}
+
+func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{maxKeys: 2}
+	auths := []*Auth{{ID: "a"}}
+
+	_, _ = selector.Pick(context.Background(), "gemini", "m1", cliproxyexecutor.Options{}, auths)
+	_, _ = selector.Pick(context.Background(), "gemini", "m2", cliproxyexecutor.Options{}, auths)
+	_, _ = selector.Pick(context.Background(), "gemini", "m3", cliproxyexecutor.Options{}, auths)
+
+	selector.mu.Lock()
+	defer selector.mu.Unlock()
+
+	if selector.cursors == nil {
+		t.Fatalf("selector.cursors = nil")
+	}
+	if len(selector.cursors) != 1 {
+		t.Fatalf("len(selector.cursors) = %d, want %d", len(selector.cursors), 1)
+	}
+	if _, ok := selector.cursors["gemini:m3"]; !ok {
+		t.Fatalf("selector.cursors missing key %q", "gemini:m3")
+	}
+}