package oauth import ( "context" "encoding/json" "errors" "fmt" "testing" "time" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/store" ) // --- mock stores --- type mockProviderStore struct { providers map[string]*store.LLMProviderData } func newMockProviderStore() *mockProviderStore { return &mockProviderStore{providers: make(map[string]*store.LLMProviderData)} } func (m *mockProviderStore) CreateProvider(_ context.Context, p *store.LLMProviderData) error { if p.ID == uuid.Nil { p.ID = uuid.New() } m.providers[p.Name] = p return nil } func (m *mockProviderStore) GetProvider(_ context.Context, id uuid.UUID) (*store.LLMProviderData, error) { for _, p := range m.providers { if p.ID == id { return p, nil } } return nil, fmt.Errorf("not found") } func (m *mockProviderStore) GetProviderByName(_ context.Context, name string) (*store.LLMProviderData, error) { if p, ok := m.providers[name]; ok { return p, nil } return nil, fmt.Errorf("not found") } func (m *mockProviderStore) ListProviders(_ context.Context) ([]store.LLMProviderData, error) { var out []store.LLMProviderData for _, p := range m.providers { out = append(out, *p) } return out, nil } func (m *mockProviderStore) UpdateProvider(_ context.Context, id uuid.UUID, updates map[string]any) error { for _, p := range m.providers { if p.ID == id { if v, ok := updates["api_key"]; ok { p.APIKey = v.(string) } if v, ok := updates["settings"]; ok { p.Settings = v.(json.RawMessage) } if v, ok := updates["enabled"]; ok { p.Enabled = v.(bool) } return nil } } return fmt.Errorf("not found") } func (m *mockProviderStore) DeleteProvider(_ context.Context, id uuid.UUID) error { for name, p := range m.providers { if p.ID == id { delete(m.providers, name) return nil } } return fmt.Errorf("not found") } func (m *mockProviderStore) ListAllProviders(_ context.Context) ([]store.LLMProviderData, error) { var out []store.LLMProviderData for _, p := range m.providers { out = append(out, *p) } return out, nil } type mockSecretsStore struct { data map[string]string } func newMockSecretsStore() *mockSecretsStore { return &mockSecretsStore{data: make(map[string]string)} } func (m *mockSecretsStore) Get(_ context.Context, key string) (string, error) { if v, ok := m.data[key]; ok { return v, nil } return "", fmt.Errorf("not found: %s", key) } func (m *mockSecretsStore) Set(_ context.Context, key, value string) error { m.data[key] = value return nil } func (m *mockSecretsStore) Delete(_ context.Context, key string) error { delete(m.data, key) return nil } func (m *mockSecretsStore) GetAll(_ context.Context) (map[string]string, error) { return m.data, nil } // --- tests --- func TestDBTokenSourceSaveAndLoad(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) resp := &OpenAITokenResponse{ AccessToken: "access-token-abc123", RefreshToken: "refresh-token-xyz789", ExpiresIn: 3600, } ctx := context.Background() id, err := ts.SaveOAuthResult(ctx, resp) if err != nil { t.Fatalf("SaveOAuthResult: %v", err) } if id == uuid.Nil { t.Fatal("expected non-nil provider ID") } // Verify provider was created p, err := provStore.GetProviderByName(ctx, DefaultProviderName) if err != nil { t.Fatalf("GetProviderByName: %v", err) } if p.APIKey != "access-token-abc123" { t.Errorf("APIKey = %q, want %q", p.APIKey, "access-token-abc123") } // Verify refresh token was saved rt, err := secretStore.Get(ctx, refreshTokenSecretKey) if err != nil { t.Fatalf("Get refresh token: %v", err) } if rt != "refresh-token-xyz789" { t.Errorf("refresh token = %q, want %q", rt, "refresh-token-xyz789") } // Load token from fresh instance ts2 := NewDBTokenSource(provStore, secretStore, DefaultProviderName) token, err := ts2.Token() if err != nil { t.Fatalf("Token: %v", err) } if token != "access-token-abc123" { t.Errorf("Token() = %q, want %q", token, "access-token-abc123") } } func TestDBTokenSourceCaching(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) resp := &OpenAITokenResponse{ AccessToken: "cached-token", RefreshToken: "refresh", ExpiresIn: 3600, } ctx := context.Background() if _, err := ts.SaveOAuthResult(ctx, resp); err != nil { t.Fatalf("SaveOAuthResult: %v", err) } // Token is cached from SaveOAuthResult token1, err := ts.Token() if err != nil { t.Fatalf("Token (1): %v", err) } // Delete provider from store — cached token should still work p, _ := provStore.GetProviderByName(ctx, DefaultProviderName) delete(provStore.providers, DefaultProviderName) token2, err := ts.Token() if err != nil { t.Fatalf("Token (2): %v", err) } if token1 != token2 { t.Errorf("cached tokens differ: %q vs %q", token1, token2) } // Restore for cleanup provStore.providers[DefaultProviderName] = p } func TestDBTokenSourceSaveOAuthResultPreservesCodexPoolSettings(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ctx := context.Background() existing := &store.LLMProviderData{ BaseModel: store.BaseModel{ID: uuid.New()}, Name: DefaultProviderName, ProviderType: store.ProviderChatGPTOAuth, APIBase: DefaultProviderAPIBase, APIKey: "old-token", Enabled: true, Settings: json.RawMessage(`{ "codex_pool": { "strategy": "round_robin", "extra_provider_names": ["openai-codex-backup"] }, "expires_at": 100 }`), } if err := provStore.CreateProvider(ctx, existing); err != nil { t.Fatalf("CreateProvider: %v", err) } ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) if _, err := ts.SaveOAuthResult(ctx, &OpenAITokenResponse{ AccessToken: "new-token", RefreshToken: "refresh-token", ExpiresIn: 3600, }); err != nil { t.Fatalf("SaveOAuthResult: %v", err) } updated, err := provStore.GetProviderByName(ctx, DefaultProviderName) if err != nil { t.Fatalf("GetProviderByName: %v", err) } settings := store.ParseChatGPTOAuthProviderSettings(updated.Settings) if settings == nil || settings.CodexPool == nil { t.Fatal("CodexPool settings lost after SaveOAuthResult") } if settings.CodexPool.Strategy != store.ChatGPTOAuthStrategyRoundRobin { t.Fatalf("Strategy = %q, want %q", settings.CodexPool.Strategy, store.ChatGPTOAuthStrategyRoundRobin) } if len(settings.CodexPool.ExtraProviderNames) != 1 || settings.CodexPool.ExtraProviderNames[0] != "openai-codex-backup" { t.Fatalf("ExtraProviderNames = %#v, want [\"openai-codex-backup\"]", settings.CodexPool.ExtraProviderNames) } } func TestDBTokenSourceExists(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) ctx := context.Background() if ts.Exists(ctx) { t.Error("Exists() = true before save, want false") } resp := &OpenAITokenResponse{ AccessToken: "token", RefreshToken: "refresh", ExpiresIn: 3600, } if _, err := ts.SaveOAuthResult(ctx, resp); err != nil { t.Fatalf("SaveOAuthResult: %v", err) } if !ts.Exists(ctx) { t.Error("Exists() = false after save, want true") } } func TestDBTokenSourceDelete(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) ctx := context.Background() resp := &OpenAITokenResponse{ AccessToken: "to-delete", RefreshToken: "refresh-to-delete", ExpiresIn: 3600, } if _, err := ts.SaveOAuthResult(ctx, resp); err != nil { t.Fatalf("SaveOAuthResult: %v", err) } if err := ts.Delete(ctx); err != nil { t.Fatalf("Delete: %v", err) } if ts.Exists(ctx) { t.Error("Exists() = true after delete, want false") } if _, err := secretStore.Get(ctx, refreshTokenSecretKey); err == nil { t.Error("refresh token still exists after delete") } } func TestDBTokenSourceExistsIgnoresNonOAuthProvider(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ctx := context.Background() if err := provStore.CreateProvider(ctx, &store.LLMProviderData{ Name: DefaultProviderName, ProviderType: store.ProviderOpenRouter, APIKey: "sk-live", Enabled: true, }); err != nil { t.Fatalf("CreateProvider: %v", err) } ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) if ts.Exists(ctx) { t.Fatal("Exists() = true, want false for non-chatgpt_oauth provider") } } func TestDBTokenSourceSaveOAuthResultRejectsProviderTypeConflict(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ctx := context.Background() if err := provStore.CreateProvider(ctx, &store.LLMProviderData{ Name: DefaultProviderName, ProviderType: store.ProviderOpenRouter, APIKey: "sk-live", Enabled: true, }); err != nil { t.Fatalf("CreateProvider: %v", err) } ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) _, err := ts.SaveOAuthResult(ctx, &OpenAITokenResponse{ AccessToken: "oauth-token", RefreshToken: "refresh-token", ExpiresIn: 3600, }) var conflict *ProviderTypeConflictError if err == nil || !errors.As(err, &conflict) { t.Fatalf("SaveOAuthResult() error = %v, want ProviderTypeConflictError", err) } } func TestDBTokenSourceUpdateExisting(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) ctx := context.Background() // Save first time resp1 := &OpenAITokenResponse{ AccessToken: "token-v1", RefreshToken: "refresh-v1", ExpiresIn: 3600, } id1, err := ts.SaveOAuthResult(ctx, resp1) if err != nil { t.Fatalf("SaveOAuthResult (1): %v", err) } // Save second time — should update, not create duplicate resp2 := &OpenAITokenResponse{ AccessToken: "token-v2", RefreshToken: "refresh-v2", ExpiresIn: 7200, } id2, err := ts.SaveOAuthResult(ctx, resp2) if err != nil { t.Fatalf("SaveOAuthResult (2): %v", err) } if id1 != id2 { t.Errorf("IDs differ on update: %s vs %s", id1, id2) } p, _ := provStore.GetProviderByName(ctx, DefaultProviderName) if p.APIKey != "token-v2" { t.Errorf("APIKey = %q after update, want %q", p.APIKey, "token-v2") } } func TestDBTokenSourceSettings(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ts := NewDBTokenSource(provStore, secretStore, DefaultProviderName) ctx := context.Background() resp := &OpenAITokenResponse{ AccessToken: "token", RefreshToken: "refresh", ExpiresIn: 3600, Scope: "openid profile", } if _, err := ts.SaveOAuthResult(ctx, resp); err != nil { t.Fatalf("SaveOAuthResult: %v", err) } p, _ := provStore.GetProviderByName(ctx, DefaultProviderName) var settings OAuthSettings if err := json.Unmarshal(p.Settings, &settings); err != nil { t.Fatalf("Unmarshal settings: %v", err) } if settings.Scopes != "openid profile" { t.Errorf("Scopes = %q, want %q", settings.Scopes, "openid profile") } expectedMin := time.Now().Add(3600 * time.Second).Unix() if settings.ExpiresAt < expectedMin-5 { t.Errorf("ExpiresAt = %d, expected >= %d", settings.ExpiresAt, expectedMin-5) } } func TestDBTokenSourceProviderScopedRefreshTokens(t *testing.T) { provStore := newMockProviderStore() secretStore := newMockSecretsStore() ctx := context.Background() first := NewDBTokenSource(provStore, secretStore, "codex-work") second := NewDBTokenSource(provStore, secretStore, "codex-personal") if _, err := first.SaveOAuthResult(ctx, &OpenAITokenResponse{ AccessToken: "token-work", RefreshToken: "refresh-work", ExpiresIn: 3600, }); err != nil { t.Fatalf("SaveOAuthResult work: %v", err) } if _, err := second.SaveOAuthResult(ctx, &OpenAITokenResponse{ AccessToken: "token-personal", RefreshToken: "refresh-personal", ExpiresIn: 3600, }); err != nil { t.Fatalf("SaveOAuthResult personal: %v", err) } workToken, err := secretStore.Get(ctx, RefreshTokenSecretKey("codex-work")) if err != nil { t.Fatalf("Get work refresh token: %v", err) } if workToken != "refresh-work" { t.Fatalf("work refresh token = %q, want %q", workToken, "refresh-work") } personalToken, err := secretStore.Get(ctx, RefreshTokenSecretKey("codex-personal")) if err != nil { t.Fatalf("Get personal refresh token: %v", err) } if personalToken != "refresh-personal" { t.Fatalf("personal refresh token = %q, want %q", personalToken, "refresh-personal") } }