package auth import ( "context" "net/http" "testing" "ds2api/internal/account" "ds2api/internal/config" ) func newTestResolver(t *testing.T) *Resolver { t.Helper() t.Setenv("DS2API_CONFIG_JSON", `{ "keys":["managed-key"], "accounts":[{"email":"acc@example.com","password":"pwd","token":"account-token"}] }`) store := config.LoadStore() pool := account.NewPool(store) return NewResolver(store, pool, func(_ context.Context, _ config.Account) (string, error) { return "fresh-token", nil }) } func TestDetermineWithXAPIKeyUsesDirectToken(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/anthropic/v1/messages", nil) req.Header.Set("x-api-key", "direct-token") auth, err := r.Determine(req) if err != nil { t.Fatalf("determine failed: %v", err) } if auth.UseConfigToken { t.Fatalf("expected direct token mode") } if auth.DeepSeekToken != "direct-token" { t.Fatalf("unexpected token: %q", auth.DeepSeekToken) } if auth.CallerID == "" { t.Fatalf("expected caller id to be populated") } } func TestDetermineWithXAPIKeyManagedKeyAcquiresAccount(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/anthropic/v1/messages", nil) req.Header.Set("x-api-key", "managed-key") auth, err := r.Determine(req) if err != nil { t.Fatalf("determine failed: %v", err) } defer r.Release(auth) if !auth.UseConfigToken { t.Fatalf("expected managed key mode") } if auth.AccountID != "acc@example.com" { t.Fatalf("unexpected account id: %q", auth.AccountID) } if auth.DeepSeekToken != "fresh-token" { t.Fatalf("unexpected account token: %q", auth.DeepSeekToken) } if auth.CallerID == "" { t.Fatalf("expected caller id to be populated") } } func TestDetermineCallerWithManagedKeySkipsAccountAcquire(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodGet, "/v1/responses/resp_1", nil) req.Header.Set("x-api-key", "managed-key") a, err := r.DetermineCaller(req) if err != nil { t.Fatalf("determine caller failed: %v", err) } if a.CallerID == "" { t.Fatalf("expected caller id to be populated") } if a.UseConfigToken { t.Fatalf("expected no config-token lease for caller-only auth") } if a.AccountID != "" { t.Fatalf("expected empty account id, got %q", a.AccountID) } } func TestCallerTokenIDStable(t *testing.T) { a := callerTokenID("token-a") b := callerTokenID("token-a") c := callerTokenID("token-b") if a == "" || b == "" || c == "" { t.Fatalf("expected non-empty caller ids") } if a != b { t.Fatalf("expected stable caller id, got %q and %q", a, b) } if a == c { t.Fatalf("expected different caller id for different tokens") } } func TestDetermineMissingToken(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil) _, err := r.Determine(req) if err == nil { t.Fatal("expected unauthorized error") } if err != ErrUnauthorized { t.Fatalf("unexpected error: %v", err) } } func TestDetermineWithQueryKeyUsesDirectToken(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent?key=direct-query-key", nil) a, err := r.Determine(req) if err != nil { t.Fatalf("determine failed: %v", err) } if a.UseConfigToken { t.Fatalf("expected direct token mode") } if a.DeepSeekToken != "direct-query-key" { t.Fatalf("unexpected token: %q", a.DeepSeekToken) } } func TestDetermineWithXGoogAPIKeyUsesDirectToken(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse", nil) req.Header.Set("x-goog-api-key", "goog-header-key") a, err := r.Determine(req) if err != nil { t.Fatalf("determine failed: %v", err) } if a.UseConfigToken { t.Fatalf("expected direct token mode") } if a.DeepSeekToken != "goog-header-key" { t.Fatalf("unexpected token: %q", a.DeepSeekToken) } } func TestDetermineWithAPIKeyQueryParamUsesDirectToken(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent?api_key=direct-api-key", nil) a, err := r.Determine(req) if err != nil { t.Fatalf("determine failed: %v", err) } if a.UseConfigToken { t.Fatalf("expected direct token mode") } if a.DeepSeekToken != "direct-api-key" { t.Fatalf("unexpected token: %q", a.DeepSeekToken) } } func TestDetermineHeaderTokenPrecedenceOverQueryKey(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent?key=query-key", nil) req.Header.Set("x-api-key", "managed-key") a, err := r.Determine(req) if err != nil { t.Fatalf("determine failed: %v", err) } defer r.Release(a) if !a.UseConfigToken { t.Fatalf("expected managed key mode from header token") } if a.AccountID == "" { t.Fatalf("expected managed account to be acquired") } } func TestDetermineCallerMissingToken(t *testing.T) { r := newTestResolver(t) req, _ := http.NewRequest(http.MethodGet, "/v1/responses/resp_1", nil) _, err := r.DetermineCaller(req) if err == nil { t.Fatal("expected unauthorized error") } if err != ErrUnauthorized { t.Fatalf("unexpected error: %v", err) } }