mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-16 07:33:06 +08:00
Erik
7f1f85cb76
Vertex AI reports 429s with a standard google.rpc.RetryInfo detail (retryDelay like "39s"), but ParseGeminiRateLimitResetTime only knew the AI Studio quotaResetDelay metadata, so parsing always failed for Vertex bodies. The fallback then lumped service_account accounts in with API-key/AI-Studio ones and rate limited them until the next PST midnight, so a transient per-minute throttle could take a pay-as-you-go account (and with it the whole group) offline for up to a day. - Parse google.rpc.RetryInfo retryDelay from 429 bodies, capped at 15 minutes to guard against absurd upstream values. - When no reset time can be parsed, cool service_account (Vertex) accounts down for 60s instead of until PST midnight: pay-as-you-go quotas are short-window RPM/TPM with no daily reset. Fallback behavior for Code Assist, Google One and API-key accounts is unchanged, covered by new regression tests.
709 lines
23 KiB
Go
709 lines
23 KiB
Go
//go:build unit
|
||
|
||
package service
|
||
|
||
import (
|
||
"context"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TestShouldFailoverGeminiUpstreamError — verifies the failover decision
|
||
// for the ErrorPolicyNone path (original logic preserved).
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestShouldFailoverGeminiUpstreamError(t *testing.T) {
|
||
svc := &GeminiMessagesCompatService{}
|
||
|
||
tests := []struct {
|
||
name string
|
||
statusCode int
|
||
expected bool
|
||
}{
|
||
{"401_failover", 401, true},
|
||
{"403_failover", 403, true},
|
||
{"429_failover", 429, true},
|
||
{"529_failover", 529, true},
|
||
{"500_failover", 500, true},
|
||
{"502_failover", 502, true},
|
||
{"503_failover", 503, true},
|
||
{"400_no_failover", 400, false},
|
||
{"404_no_failover", 404, false},
|
||
{"422_no_failover", 422, false},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
got := svc.shouldFailoverGeminiUpstreamError(tt.statusCode)
|
||
require.Equal(t, tt.expected, got)
|
||
})
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TestCheckErrorPolicy_GeminiAccounts — verifies CheckErrorPolicy works
|
||
// correctly for Gemini platform accounts (API Key type).
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestCheckErrorPolicy_GeminiAccounts(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
account *Account
|
||
statusCode int
|
||
body []byte
|
||
expected ErrorPolicyResult
|
||
}{
|
||
{
|
||
name: "gemini_apikey_custom_codes_hit",
|
||
account: &Account{
|
||
ID: 100,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429), float64(500)},
|
||
},
|
||
},
|
||
statusCode: 429,
|
||
body: []byte(`{"error":"rate limited"}`),
|
||
expected: ErrorPolicyMatched,
|
||
},
|
||
{
|
||
name: "gemini_apikey_custom_codes_miss",
|
||
account: &Account{
|
||
ID: 101,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
},
|
||
statusCode: 500,
|
||
body: []byte(`{"error":"internal"}`),
|
||
expected: ErrorPolicySkipped,
|
||
},
|
||
{
|
||
name: "gemini_apikey_no_custom_codes_returns_none",
|
||
account: &Account{
|
||
ID: 102,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
},
|
||
statusCode: 500,
|
||
body: []byte(`{"error":"internal"}`),
|
||
expected: ErrorPolicyNone,
|
||
},
|
||
{
|
||
name: "gemini_apikey_temp_unschedulable_hit",
|
||
account: &Account{
|
||
ID: 103,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"temp_unschedulable_enabled": true,
|
||
"temp_unschedulable_rules": []any{
|
||
map[string]any{
|
||
"error_code": float64(503),
|
||
"keywords": []any{"overloaded"},
|
||
"duration_minutes": float64(10),
|
||
},
|
||
},
|
||
},
|
||
},
|
||
statusCode: 503,
|
||
body: []byte(`overloaded service`),
|
||
expected: ErrorPolicyTempUnscheduled,
|
||
},
|
||
{
|
||
name: "gemini_apikey_temp_unschedulable_401_second_hit_returns_none",
|
||
account: &Account{
|
||
ID: 105,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
TempUnschedulableReason: `{"status_code":401,"until_unix":1735689600}`,
|
||
Credentials: map[string]any{
|
||
"temp_unschedulable_enabled": true,
|
||
"temp_unschedulable_rules": []any{
|
||
map[string]any{
|
||
"error_code": float64(401),
|
||
"keywords": []any{"unauthorized"},
|
||
"duration_minutes": float64(10),
|
||
},
|
||
},
|
||
},
|
||
},
|
||
statusCode: 401,
|
||
body: []byte(`unauthorized`),
|
||
expected: ErrorPolicyNone,
|
||
},
|
||
{
|
||
name: "gemini_custom_codes_override_temp_unschedulable",
|
||
account: &Account{
|
||
ID: 104,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(503)},
|
||
"temp_unschedulable_enabled": true,
|
||
"temp_unschedulable_rules": []any{
|
||
map[string]any{
|
||
"error_code": float64(503),
|
||
"keywords": []any{"overloaded"},
|
||
"duration_minutes": float64(10),
|
||
},
|
||
},
|
||
},
|
||
},
|
||
statusCode: 503,
|
||
body: []byte(`overloaded`),
|
||
expected: ErrorPolicyMatched, // custom codes take precedence
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
repo := &errorPolicyRepoStub{}
|
||
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||
|
||
result := svc.CheckErrorPolicy(context.Background(), tt.account, tt.statusCode, tt.body)
|
||
require.Equal(t, tt.expected, result)
|
||
})
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TestGeminiErrorPolicyIntegration — verifies the Gemini error handling
|
||
// paths produce the correct behavior for each ErrorPolicyResult.
|
||
//
|
||
// These tests simulate the inline error policy switch in handleClaudeCompat
|
||
// and forwardNativeGemini by calling the same methods in the same order.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestGeminiErrorPolicyIntegration(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
tests := []struct {
|
||
name string
|
||
account *Account
|
||
statusCode int
|
||
respBody []byte
|
||
expectFailover bool // expect UpstreamFailoverError
|
||
expectHandleError bool // expect handleGeminiUpstreamError to be called
|
||
expectShouldFailover bool // for None path, whether shouldFailover triggers
|
||
expectModelScope string
|
||
}{
|
||
{
|
||
name: "custom_codes_matched_429_failover",
|
||
account: &Account{
|
||
ID: 200,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
},
|
||
statusCode: 429,
|
||
respBody: []byte(`{"error":"rate limited"}`),
|
||
expectFailover: true,
|
||
expectHandleError: true,
|
||
},
|
||
{
|
||
name: "custom_codes_skipped_500_failover",
|
||
account: &Account{
|
||
ID: 201,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
},
|
||
statusCode: 500,
|
||
respBody: []byte(`{"error":"internal"}`),
|
||
expectFailover: true,
|
||
expectHandleError: false,
|
||
},
|
||
{
|
||
name: "custom_codes_skipped_400_no_failover",
|
||
account: &Account{
|
||
ID: 205,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
},
|
||
statusCode: 400,
|
||
respBody: []byte(`{"error":"bad request"}`),
|
||
expectFailover: false,
|
||
expectHandleError: false,
|
||
},
|
||
{
|
||
name: "temp_unschedulable_matched_failover",
|
||
account: &Account{
|
||
ID: 202,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"temp_unschedulable_enabled": true,
|
||
"temp_unschedulable_rules": []any{
|
||
map[string]any{
|
||
"error_code": float64(503),
|
||
"keywords": []any{"overloaded"},
|
||
"duration_minutes": float64(10),
|
||
},
|
||
},
|
||
},
|
||
},
|
||
statusCode: 503,
|
||
respBody: []byte(`overloaded`),
|
||
expectFailover: true,
|
||
expectHandleError: false,
|
||
expectModelScope: "gemini-2.5-pro",
|
||
},
|
||
{
|
||
name: "no_policy_429_failover_via_shouldFailover",
|
||
account: &Account{
|
||
ID: 203,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
},
|
||
statusCode: 429,
|
||
respBody: []byte(`{"error":"rate limited"}`),
|
||
expectFailover: true,
|
||
expectHandleError: true,
|
||
expectShouldFailover: true,
|
||
},
|
||
{
|
||
name: "no_policy_400_no_failover",
|
||
account: &Account{
|
||
ID: 204,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
},
|
||
statusCode: 400,
|
||
respBody: []byte(`{"error":"bad request"}`),
|
||
expectFailover: false,
|
||
expectHandleError: true,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
repo := &geminiErrorPolicyRepo{}
|
||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||
svc := &GeminiMessagesCompatService{
|
||
accountRepo: repo,
|
||
rateLimitService: rlSvc,
|
||
}
|
||
|
||
writer := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(writer)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
// Simulate the Claude compat error handling path (same logic as native).
|
||
// This mirrors the inline switch in handleClaudeCompat.
|
||
var handleErrorCalled bool
|
||
var gotFailover bool
|
||
|
||
ctx := context.Background()
|
||
statusCode := tt.statusCode
|
||
respBody := tt.respBody
|
||
account := tt.account
|
||
headers := http.Header{}
|
||
|
||
if svc.rateLimitService != nil {
|
||
policy := svc.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, respBody, "gemini-2.5-pro")
|
||
switch policy {
|
||
case ErrorPolicySkipped:
|
||
// Skipped → 不标记账号状态；可 failover 的状态码仍换号
|
||
handleErrorCalled = false
|
||
gotFailover = svc.skippedErrorPolicyFailoverError(c, account, statusCode, respBody, "req-test") != nil
|
||
goto verify
|
||
case ErrorPolicyMatched:
|
||
svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody)
|
||
handleErrorCalled = true
|
||
gotFailover = true
|
||
goto verify
|
||
case ErrorPolicyTempUnscheduled:
|
||
handleErrorCalled = false
|
||
gotFailover = true
|
||
goto verify
|
||
}
|
||
}
|
||
|
||
// ErrorPolicyNone → original logic
|
||
svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody)
|
||
handleErrorCalled = true
|
||
if svc.shouldFailoverGeminiUpstreamError(statusCode) {
|
||
gotFailover = true
|
||
}
|
||
|
||
verify:
|
||
require.Equal(t, tt.expectFailover, gotFailover, "failover mismatch")
|
||
require.Equal(t, tt.expectHandleError, handleErrorCalled, "handleGeminiUpstreamError call mismatch")
|
||
if tt.expectModelScope != "" {
|
||
require.Equal(t, 1, repo.setModelRateLimitedCalls)
|
||
require.Equal(t, tt.expectModelScope, repo.lastModelScope)
|
||
require.Zero(t, repo.setTempCalls)
|
||
require.Zero(t, repo.setRateLimitedCalls, "model temp rule must not be widened into an account rate limit")
|
||
}
|
||
|
||
if tt.expectShouldFailover {
|
||
require.True(t, svc.shouldFailoverGeminiUpstreamError(statusCode),
|
||
"shouldFailoverGeminiUpstreamError should return true for status %d", statusCode)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TestSkippedErrorPolicyFailoverError — ErrorPolicySkipped（池模式、或自定义
|
||
// 错误码未命中）不豁免换号：可 failover 的状态码返回 UpstreamFailoverError，
|
||
// 仅池模式账号可携带同账号重试标记。
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestSkippedErrorPolicyFailoverError(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
svc := &GeminiMessagesCompatService{}
|
||
|
||
poolAccount := func(extra map[string]any) *Account {
|
||
creds := map[string]any{"pool_mode": true}
|
||
for k, v := range extra {
|
||
creds[k] = v
|
||
}
|
||
return &Account{ID: 300, Type: AccountTypeAPIKey, Platform: PlatformGemini, Credentials: creds}
|
||
}
|
||
customCodesAccount := &Account{
|
||
ID: 301, Type: AccountTypeAPIKey, Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
}
|
||
|
||
tests := []struct {
|
||
name string
|
||
account *Account
|
||
statusCode int
|
||
expectFailover bool
|
||
expectSameAccount bool
|
||
}{
|
||
{"pool_500_failover_no_same_account_retry", poolAccount(nil), 500, true, false},
|
||
{"pool_429_failover_with_same_account_retry", poolAccount(nil), 429, true, true},
|
||
{"pool_custom_retry_codes_500", poolAccount(map[string]any{
|
||
"pool_mode_retry_status_codes": []any{float64(500)},
|
||
}), 500, true, true},
|
||
{"pool_400_not_failover_worthy", poolAccount(nil), 400, false, false},
|
||
{"custom_codes_miss_500_failover_no_same_account_retry", customCodesAccount, 500, true, false},
|
||
{"custom_codes_miss_400_not_failover_worthy", customCodesAccount, 400, false, false},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
writer := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(writer)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||
|
||
body := []byte(`{"error":{"code":"bad_response_status_code","message":"openai_error"}}`)
|
||
failoverErr := svc.skippedErrorPolicyFailoverError(c, tt.account, tt.statusCode, body, "req-1")
|
||
|
||
if !tt.expectFailover {
|
||
require.Nil(t, failoverErr)
|
||
return
|
||
}
|
||
require.NotNil(t, failoverErr)
|
||
require.Equal(t, tt.statusCode, failoverErr.StatusCode)
|
||
require.Equal(t, body, failoverErr.ResponseBody)
|
||
require.Equal(t, tt.expectSameAccount, failoverErr.RetryableOnSameAccount)
|
||
require.True(t, failoverErr.ShouldRetryNextAccount())
|
||
})
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TestGeminiErrorPolicy_NilRateLimitService — verifies nil safety
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestGeminiErrorPolicy_NilRateLimitService(t *testing.T) {
|
||
svc := &GeminiMessagesCompatService{
|
||
rateLimitService: nil,
|
||
}
|
||
|
||
// When rateLimitService is nil, error policy is skipped → falls through to
|
||
// shouldFailoverGeminiUpstreamError (original logic).
|
||
// Verify this doesn't panic and follows expected behavior.
|
||
|
||
ctx := context.Background()
|
||
account := &Account{
|
||
ID: 300,
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformGemini,
|
||
Credentials: map[string]any{
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
}
|
||
|
||
// The nil check should prevent CheckErrorPolicy from being called
|
||
if svc.rateLimitService != nil {
|
||
t.Fatal("rateLimitService should be nil for this test")
|
||
}
|
||
|
||
// shouldFailoverGeminiUpstreamError still works
|
||
require.True(t, svc.shouldFailoverGeminiUpstreamError(429))
|
||
require.False(t, svc.shouldFailoverGeminiUpstreamError(400))
|
||
|
||
// handleGeminiUpstreamError should not panic with nil rateLimitService
|
||
require.NotPanics(t, func() {
|
||
svc.handleGeminiUpstreamError(ctx, account, 500, http.Header{}, []byte(`error`))
|
||
})
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// geminiErrorPolicyRepo — minimal AccountRepository stub for Gemini error
|
||
// policy tests. Embeds mockAccountRepoForGemini and adds tracking.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestHandleGeminiUpstreamError_GoogleOneCapacityExhaustedUsesTierCooldown(t *testing.T) {
|
||
repo := &rateLimit429AccountRepoStub{}
|
||
quotaSvc := NewGeminiQuotaService(&config.Config{}, nil)
|
||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, quotaSvc, nil)
|
||
svc := &GeminiMessagesCompatService{
|
||
accountRepo: repo,
|
||
rateLimitService: rlSvc,
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 511,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{
|
||
"oauth_type": "google_one",
|
||
"tier_id": "google_ai_pro",
|
||
},
|
||
}
|
||
body := []byte(`{"error":{"code":429,"details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","domain":"cloudcode-pa.googleapis.com","metadata":{"model":"gemini-3.1-pro-preview"},"reason":"MODEL_CAPACITY_EXHAUSTED"}],"message":"No capacity available for model gemini-3.1-pro-preview on the server","status":"RESOURCE_EXHAUSTED"}}`)
|
||
|
||
before := time.Now()
|
||
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
|
||
after := time.Now()
|
||
|
||
require.Equal(t, 1, repo.rateLimitCalls)
|
||
require.Equal(t, int64(511), repo.lastRateLimitID)
|
||
require.WithinDuration(t, before.Add(5*time.Minute), repo.lastRateLimitReset, 2*time.Second)
|
||
require.True(t, repo.lastRateLimitReset.After(before))
|
||
require.True(t, repo.lastRateLimitReset.Before(after.Add(5*time.Minute).Add(2*time.Second)))
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// TestHandleGeminiUpstreamError_PoolMode429 — 池模式账号的 429 不写账号级限流。
|
||
//
|
||
// 429 的标记点在重试循环内（handleClaudeCompat / forwardNativeGemini /
|
||
// chat completions 三条路径），先于 CheckErrorPolicy 执行，池模式豁免只能落在
|
||
// handleGeminiUpstreamError 自身；否则一次上游 429 会把账号锁到 PST 午夜，
|
||
// 即便重试已经成功返回客户端。
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestHandleGeminiUpstreamError_PoolMode429(t *testing.T) {
|
||
// 中转上游的真实 429 文案：不含 "per day"，也没有 quotaResetDelay，
|
||
// 解析失败后 apikey 账号会落到 PST 午夜兜底。
|
||
body := []byte(`{"error":{"code":429,"message":"You have exhausted your capacity on this model. Your quota will reset after 6h53m10s."}}`)
|
||
|
||
tests := []struct {
|
||
name string
|
||
account *Account
|
||
expectRateLimited bool
|
||
}{
|
||
{
|
||
name: "pool_mode_apikey_stays_in_pool",
|
||
account: &Account{
|
||
ID: 600,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{"pool_mode": true},
|
||
},
|
||
expectRateLimited: false,
|
||
},
|
||
{
|
||
name: "custom_error_codes_hit_overrides_pool_mode",
|
||
account: &Account{
|
||
ID: 601,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{
|
||
"pool_mode": true,
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(429)},
|
||
},
|
||
},
|
||
expectRateLimited: true,
|
||
},
|
||
{
|
||
name: "custom_error_codes_miss_skips",
|
||
account: &Account{
|
||
ID: 602,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{
|
||
"pool_mode": true,
|
||
"custom_error_codes_enabled": true,
|
||
"custom_error_codes": []any{float64(500)},
|
||
},
|
||
},
|
||
expectRateLimited: false,
|
||
},
|
||
{
|
||
name: "non_pool_apikey_still_rate_limited",
|
||
account: &Account{
|
||
ID: 603,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeAPIKey,
|
||
},
|
||
expectRateLimited: true,
|
||
},
|
||
{
|
||
name: "oauth_account_ignores_pool_mode_flag",
|
||
account: &Account{
|
||
ID: 604,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"pool_mode": true},
|
||
},
|
||
expectRateLimited: true,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
repo := &rateLimit429AccountRepoStub{}
|
||
svc := &GeminiMessagesCompatService{
|
||
accountRepo: repo,
|
||
rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
|
||
}
|
||
|
||
svc.handleGeminiUpstreamError(context.Background(), tt.account, http.StatusTooManyRequests, http.Header{}, body)
|
||
|
||
if !tt.expectRateLimited {
|
||
require.Zero(t, repo.rateLimitCalls, "池模式账号不应被标记账号级限流")
|
||
return
|
||
}
|
||
require.Equal(t, 1, repo.rateLimitCalls)
|
||
require.Equal(t, tt.account.ID, repo.lastRateLimitID)
|
||
require.True(t, repo.lastRateLimitReset.After(time.Now()))
|
||
})
|
||
}
|
||
}
|
||
|
||
// Vertex（service_account）429 响应带 google.rpc.RetryInfo 时，应按 retryDelay 冷却，
|
||
// 而不是回退到 PST 午夜。
|
||
func TestHandleGeminiUpstreamError_VertexRetryInfoUsesParsedDelay(t *testing.T) {
|
||
repo := &rateLimit429AccountRepoStub{}
|
||
svc := &GeminiMessagesCompatService{accountRepo: repo}
|
||
|
||
account := &Account{
|
||
ID: 601,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeServiceAccount,
|
||
Credentials: map[string]any{
|
||
"project_id": "my-vertex-project",
|
||
"location": "global",
|
||
},
|
||
}
|
||
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"39s"}]}}`)
|
||
|
||
before := time.Now()
|
||
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
|
||
|
||
require.Equal(t, 1, repo.rateLimitCalls)
|
||
require.Equal(t, int64(601), repo.lastRateLimitID)
|
||
require.WithinDuration(t, before.Add(39*time.Second), repo.lastRateLimitReset, 2*time.Second)
|
||
}
|
||
|
||
// Vertex（service_account）429 响应不带任何可解析的重置时间时，走短冷却兜底，
|
||
// 不再冷却到 PST 午夜（按量付费没有每日配额，长冷却会导致整组不可用）。
|
||
func TestHandleGeminiUpstreamError_VertexFallbackUsesShortCooldown(t *testing.T) {
|
||
repo := &rateLimit429AccountRepoStub{}
|
||
svc := &GeminiMessagesCompatService{accountRepo: repo}
|
||
|
||
account := &Account{
|
||
ID: 602,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeServiceAccount,
|
||
Credentials: map[string]any{
|
||
"project_id": "my-vertex-project",
|
||
},
|
||
}
|
||
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED"}}`)
|
||
|
||
before := time.Now()
|
||
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
|
||
|
||
require.Equal(t, 1, repo.rateLimitCalls)
|
||
require.Equal(t, int64(602), repo.lastRateLimitID)
|
||
require.WithinDuration(t, before.Add(geminiVertexFallbackCooldown), repo.lastRateLimitReset, 2*time.Second)
|
||
// 显式守护：不应被冷却超过 15 分钟（更不应到次日 PST 午夜）
|
||
require.True(t, repo.lastRateLimitReset.Before(before.Add(15*time.Minute)))
|
||
}
|
||
|
||
// API Key（AI Studio）无法解析重置时间时仍冷却到 PST 午夜，行为保持不变。
|
||
func TestHandleGeminiUpstreamError_APIKeyFallbackStillPSTMidnight(t *testing.T) {
|
||
repo := &rateLimit429AccountRepoStub{}
|
||
svc := &GeminiMessagesCompatService{accountRepo: repo}
|
||
|
||
account := &Account{
|
||
ID: 603,
|
||
Platform: PlatformGemini,
|
||
Type: AccountTypeAPIKey,
|
||
}
|
||
body := []byte(`{"error":{"code":429,"message":"rate limit"}}`)
|
||
|
||
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
|
||
|
||
require.Equal(t, 1, repo.rateLimitCalls)
|
||
require.Equal(t, int64(603), repo.lastRateLimitID)
|
||
expected := time.Unix(*nextGeminiDailyResetUnix(), 0)
|
||
require.WithinDuration(t, expected, repo.lastRateLimitReset, 5*time.Second)
|
||
}
|
||
|
||
type geminiErrorPolicyRepo struct {
|
||
mockAccountRepoForGemini
|
||
setErrorCalls int
|
||
setRateLimitedCalls int
|
||
setTempCalls int
|
||
setModelRateLimitedCalls int
|
||
lastModelScope string
|
||
}
|
||
|
||
func (r *geminiErrorPolicyRepo) SetError(_ context.Context, _ int64, _ string) error {
|
||
r.setErrorCalls++
|
||
return nil
|
||
}
|
||
|
||
func (r *geminiErrorPolicyRepo) SetRateLimited(_ context.Context, _ int64, _ time.Time) error {
|
||
r.setRateLimitedCalls++
|
||
return nil
|
||
}
|
||
|
||
func (r *geminiErrorPolicyRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, _ string) error {
|
||
r.setTempCalls++
|
||
return nil
|
||
}
|
||
|
||
func (r *geminiErrorPolicyRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, _ time.Time, _ ...string) error {
|
||
r.setModelRateLimitedCalls++
|
||
r.lastModelScope = scope
|
||
return nil
|
||
}
|