Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/token_refresh_service_candidates_test.go

@@ -0,0 +1,246 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/imroc/req/v3"
+	"github.com/stretchr/testify/require"
+)
+
+type tokenRefreshCandidateRepo struct {
+	AccountRepository
+	mu                    sync.Mutex
+	accounts              []Account
+	updatedCredentialIDs  []int64
+	setErrorCalls         int
+	setTempUnschedCalls   int
+	clearTempCalls        int
+	lastTempUnschedReason string
+	listActiveCalls       int
+}
+
+func (r *tokenRefreshCandidateRepo) ListActive(context.Context) ([]Account, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.listActiveCalls++
+	return r.accounts, nil
+}
+
+func (r *tokenRefreshCandidateRepo) ListOAuthRefreshCandidatePage(_ context.Context, options OAuthRefreshPageOptions) (*OAuthRefreshCandidatePage, error) {
+	candidates := make([]Account, 0, len(r.accounts))
+	now := time.Now()
+	for _, account := range r.accounts {
+		if account.ID <= options.AfterID {
+			continue
+		}
+		refreshToken, _ := account.Credentials["refresh_token"].(string)
+		inRetryCooldown := account.TempUnschedulableUntil != nil &&
+			account.TempUnschedulableUntil.After(now) &&
+			strings.HasPrefix(account.TempUnschedulableReason, "token refresh retry exhausted:")
+		platformAllowed := false
+		for _, platform := range options.Platforms {
+			if account.Platform == platform {
+				platformAllowed = true
+				break
+			}
+		}
+		if options.ActiveOnly && account.Status != StatusActive ||
+			!account.Schedulable ||
+			account.Type != AccountTypeOAuth ||
+			!platformAllowed ||
+			options.RequireRefreshToken && strings.TrimSpace(refreshToken) == "" ||
+			options.ExcludeRetryCooldown && inRetryCooldown {
+			continue
+		}
+		candidates = append(candidates, account)
+		if len(candidates) == options.Limit {
+			break
+		}
+	}
+	page := &OAuthRefreshCandidatePage{Accounts: candidates, HasMore: len(candidates) == options.Limit}
+	if len(candidates) > 0 {
+		page.NextAfterID = candidates[len(candidates)-1].ID
+	}
+	return page, nil
+}
+
+func (r *tokenRefreshCandidateRepo) UpdateCredentials(_ context.Context, id int64, _ map[string]any) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
+	return nil
+}
+
+func (r *tokenRefreshCandidateRepo) SetError(context.Context, int64, string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.setErrorCalls++
+	return nil
+}
+
+func (r *tokenRefreshCandidateRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.setTempUnschedCalls++
+	r.lastTempUnschedReason = reason
+	return nil
+}
+
+func (r *tokenRefreshCandidateRepo) ClearTempUnschedulable(context.Context, int64) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.clearTempCalls++
+	return nil
+}
+
+type tokenRefreshTestRefresher struct {
+	err error
+}
+
+func (r *tokenRefreshTestRefresher) CanRefresh(*Account) bool { return true }
+
+func (r *tokenRefreshTestRefresher) NeedsRefresh(*Account, time.Duration) bool { return true }
+
+func (r *tokenRefreshTestRefresher) Refresh(context.Context, *Account) (map[string]any, error) {
+	if r.err != nil {
+		return nil, r.err
+	}
+	return map[string]any{"access_token": "new-access-token", "refresh_token": "new-refresh-token"}, nil
+}
+
+func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing.T) {
+	future := time.Now().Add(10 * time.Minute)
+	repo := &tokenRefreshCandidateRepo{
+		accounts: []Account{
+			{
+				ID:          1,
+				Platform:    PlatformOpenAI,
+				Type:        AccountTypeOAuth,
+				Status:      StatusActive,
+				Schedulable: true,
+				Credentials: map[string]any{"refresh_token": "refresh-token"},
+			},
+			{
+				ID:          2,
+				Platform:    PlatformOpenAI,
+				Type:        AccountTypeOAuth,
+				Status:      StatusActive,
+				Schedulable: true,
+				Credentials: map[string]any{},
+			},
+			{
+				ID:          3,
+				Platform:    PlatformGemini,
+				Type:        AccountTypeAPIKey,
+				Status:      StatusActive,
+				Credentials: map[string]any{"refresh_token": "refresh-token"},
+			},
+			{
+				ID:                      4,
+				Platform:                PlatformAntigravity,
+				Type:                    AccountTypeOAuth,
+				Status:                  StatusActive,
+				Schedulable:             true,
+				Credentials:             map[string]any{"refresh_token": "refresh-token"},
+				TempUnschedulableUntil:  &future,
+				TempUnschedulableReason: "token refresh retry exhausted: network timeout",
+			},
+			{
+				ID:          5,
+				Platform:    "other",
+				Type:        AccountTypeOAuth,
+				Status:      StatusActive,
+				Schedulable: true,
+				Credentials: map[string]any{"refresh_token": "refresh-token"},
+			},
+			{
+				ID:                      6,
+				Platform:                PlatformAntigravity,
+				Type:                    AccountTypeOAuth,
+				Status:                  StatusActive,
+				Schedulable:             true,
+				Credentials:             map[string]any{"refresh_token": "refresh-token"},
+				Extra:                   map[string]any{"privacy_mode": AntigravityPrivacySet},
+				TempUnschedulableUntil:  &future,
+				TempUnschedulableReason: "OAuth 401: unauthorized",
+			},
+			{
+				ID:          7,
+				Platform:    PlatformOpenAI,
+				Type:        AccountTypeOAuth,
+				Status:      StatusActive,
+				Schedulable: false,
+				Credentials: map[string]any{"refresh_token": "permanently-rejected-token"},
+			},
+		},
+	}
+	svc := &TokenRefreshService{
+		accountRepo:    repo,
+		candidatePager: repo,
+		registrations: []tokenRefreshRegistration{
+			{platform: PlatformOpenAI, refresher: &tokenRefreshTestRefresher{}},
+			{platform: PlatformGemini, refresher: &tokenRefreshTestRefresher{}},
+			{platform: PlatformAntigravity, refresher: &tokenRefreshTestRefresher{}},
+		},
+		refreshPolicy: DefaultBackgroundRefreshPolicy(),
+		cfg:           &config.TokenRefreshConfig{RefreshBeforeExpiryHours: 1, MaxRetries: 1},
+	}
+
+	svc.processRefresh()
+
+	require.Zero(t, repo.listActiveCalls, "TokenRefreshService should not use the broad active-account query")
+	require.ElementsMatch(t, []int64{1, 6}, repo.updatedCredentialIDs)
+	require.Equal(t, 1, repo.clearTempCalls, "successful refresh should clear the OAuth 401 temp-unschedulable state")
+}
+
+func TestTokenRefreshService_RefreshFailureDoesNotCallPrivacy(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+	}{
+		{name: "retry exhausted", err: errors.New("temporary upstream timeout")},
+		{name: "non retryable", err: errors.New("invalid_grant: token revoked")},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			repo := &tokenRefreshCandidateRepo{}
+			svc := &TokenRefreshService{
+				accountRepo:   repo,
+				refreshPolicy: DefaultBackgroundRefreshPolicy(),
+				cfg:           &config.TokenRefreshConfig{MaxRetries: 1, RetryBackoffSeconds: 0},
+				privacyClientFactory: func(string) (*req.Client, error) {
+					t.Fatalf("privacy client factory must not be called on refresh failure")
+					return nil, errors.New("unexpected privacy call")
+				},
+			}
+			account := &Account{
+				ID:       11,
+				Platform: PlatformOpenAI,
+				Type:     AccountTypeOAuth,
+				Credentials: map[string]any{
+					"access_token":  "old-access-token",
+					"refresh_token": "refresh-token",
+				},
+			}
+
+			err := svc.refreshWithRetry(context.Background(), account, &tokenRefreshTestRefresher{err: tt.err}, nil, time.Hour)
+
+			require.Error(t, err)
+			if isNonRetryableRefreshError(tt.err) {
+				require.Equal(t, 1, repo.setErrorCalls)
+				require.Zero(t, repo.setTempUnschedCalls)
+			} else {
+				require.Zero(t, repo.setErrorCalls)
+				require.Equal(t, 1, repo.setTempUnschedCalls)
+				require.True(t, strings.HasPrefix(repo.lastTempUnschedReason, "token refresh retry exhausted:"))
+			}
+		})
+	}
+}