Merge branch 'main' into test
This commit is contained in:
@@ -92,7 +92,9 @@ func (s *stubAntigravityAccountRepo) SetModelRateLimit(ctx context.Context, id i
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
|
||||
func TestAntigravityRetryLoop_NoURLFallback_UsesConfiguredBaseURL(t *testing.T) {
|
||||
t.Setenv(antigravityForwardBaseURLEnv, "")
|
||||
|
||||
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
|
||||
oldAvailability := antigravity.DefaultURLAvailability
|
||||
defer func() {
|
||||
@@ -137,15 +139,16 @@ func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.resp)
|
||||
defer func() { _ = result.resp.Body.Close() }()
|
||||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||||
require.False(t, handleErrorCalled)
|
||||
require.Len(t, upstream.calls, 2)
|
||||
require.True(t, strings.HasPrefix(upstream.calls[0], base1))
|
||||
require.True(t, strings.HasPrefix(upstream.calls[1], base2))
|
||||
require.Equal(t, http.StatusTooManyRequests, result.resp.StatusCode)
|
||||
require.True(t, handleErrorCalled)
|
||||
require.Len(t, upstream.calls, antigravityMaxRetries)
|
||||
for _, callURL := range upstream.calls {
|
||||
require.True(t, strings.HasPrefix(callURL, base1))
|
||||
}
|
||||
|
||||
available := antigravity.DefaultURLAvailability.GetAvailableURLs()
|
||||
require.NotEmpty(t, available)
|
||||
require.Equal(t, base2, available[0])
|
||||
require.Equal(t, base1, available[0])
|
||||
}
|
||||
|
||||
// TestHandleUpstreamError_429_ModelRateLimit 测试 429 模型限流场景
|
||||
@@ -194,13 +197,14 @@ func TestHandleUpstreamError_429_NonModelRateLimit(t *testing.T) {
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// TestHandleUpstreamError_503_ModelRateLimit 测试 503 模型限流场景
|
||||
func TestHandleUpstreamError_503_ModelRateLimit(t *testing.T) {
|
||||
// TestHandleUpstreamError_503_ModelCapacityExhausted 测试 503 模型容量不足场景
|
||||
// MODEL_CAPACITY_EXHAUSTED 时应等待重试,不切换账号
|
||||
func TestHandleUpstreamError_503_ModelCapacityExhausted(t *testing.T) {
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
svc := &AntigravityGatewayService{accountRepo: repo}
|
||||
account := &Account{ID: 3, Name: "acc-3", Platform: PlatformAntigravity}
|
||||
|
||||
// 503 + MODEL_CAPACITY_EXHAUSTED → 模型限流
|
||||
// 503 + MODEL_CAPACITY_EXHAUSTED → 等待重试,不切换账号
|
||||
body := []byte(`{
|
||||
"error": {
|
||||
"status": "UNAVAILABLE",
|
||||
@@ -213,13 +217,13 @@ func TestHandleUpstreamError_503_ModelRateLimit(t *testing.T) {
|
||||
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
|
||||
|
||||
// 应该触发模型限流
|
||||
// MODEL_CAPACITY_EXHAUSTED 应该标记为已处理,不切换账号,不设置模型限流
|
||||
// 实际重试由 handleSmartRetry 处理
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.Handled)
|
||||
require.NotNil(t, result.SwitchError)
|
||||
require.Equal(t, "gemini-3-pro-high", result.SwitchError.RateLimitedModel)
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "gemini-3-pro-high", repo.modelRateLimitCalls[0].modelKey)
|
||||
require.False(t, result.ShouldRetry, "MODEL_CAPACITY_EXHAUSTED should not trigger retry from handleModelRateLimit path")
|
||||
require.Nil(t, result.SwitchError, "MODEL_CAPACITY_EXHAUSTED should not trigger account switch")
|
||||
require.Empty(t, repo.modelRateLimitCalls, "MODEL_CAPACITY_EXHAUSTED should not set model rate limit")
|
||||
}
|
||||
|
||||
// TestHandleUpstreamError_503_NonModelRateLimit 测试 503 非模型限流场景(不处理)
|
||||
@@ -307,11 +311,12 @@ func TestParseGeminiRateLimitResetTime_QuotaResetDelay_RoundsUp(t *testing.T) {
|
||||
|
||||
func TestParseAntigravitySmartRetryInfo(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
expectedDelay time.Duration
|
||||
expectedModel string
|
||||
expectedNil bool
|
||||
name string
|
||||
body string
|
||||
expectedDelay time.Duration
|
||||
expectedModel string
|
||||
expectedNil bool
|
||||
expectedIsModelCapacityExhausted bool
|
||||
}{
|
||||
{
|
||||
name: "valid complete response with RATE_LIMIT_EXCEEDED",
|
||||
@@ -374,8 +379,9 @@ func TestParseAntigravitySmartRetryInfo(t *testing.T) {
|
||||
"message": "No capacity available for model gemini-3-pro-high on the server"
|
||||
}
|
||||
}`,
|
||||
expectedDelay: 39 * time.Second,
|
||||
expectedModel: "gemini-3-pro-high",
|
||||
expectedDelay: 39 * time.Second,
|
||||
expectedModel: "gemini-3-pro-high",
|
||||
expectedIsModelCapacityExhausted: true,
|
||||
},
|
||||
{
|
||||
name: "503 UNAVAILABLE without MODEL_CAPACITY_EXHAUSTED - should return nil",
|
||||
@@ -486,6 +492,9 @@ func TestParseAntigravitySmartRetryInfo(t *testing.T) {
|
||||
if result.ModelName != tt.expectedModel {
|
||||
t.Errorf("ModelName = %q, want %q", result.ModelName, tt.expectedModel)
|
||||
}
|
||||
if result.IsModelCapacityExhausted != tt.expectedIsModelCapacityExhausted {
|
||||
t.Errorf("IsModelCapacityExhausted = %v, want %v", result.IsModelCapacityExhausted, tt.expectedIsModelCapacityExhausted)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -497,13 +506,14 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
apiKeyAccount := &Account{Type: AccountTypeAPIKey}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
body string
|
||||
expectedShouldRetry bool
|
||||
expectedShouldRateLimit bool
|
||||
minWait time.Duration
|
||||
modelName string
|
||||
name string
|
||||
account *Account
|
||||
body string
|
||||
expectedShouldRetry bool
|
||||
expectedShouldRateLimit bool
|
||||
expectedIsModelCapacityExhausted bool
|
||||
minWait time.Duration
|
||||
modelName string
|
||||
}{
|
||||
{
|
||||
name: "OAuth account with short delay (< 7s) - smart retry",
|
||||
@@ -617,13 +627,14 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
]
|
||||
}
|
||||
}`,
|
||||
expectedShouldRetry: false,
|
||||
expectedShouldRateLimit: true,
|
||||
minWait: 39 * time.Second,
|
||||
modelName: "gemini-3-pro-high",
|
||||
expectedShouldRetry: true,
|
||||
expectedShouldRateLimit: false,
|
||||
expectedIsModelCapacityExhausted: true,
|
||||
minWait: 1 * time.Second,
|
||||
modelName: "gemini-3-pro-high",
|
||||
},
|
||||
{
|
||||
name: "503 UNAVAILABLE with MODEL_CAPACITY_EXHAUSTED - no retryDelay - use default rate limit",
|
||||
name: "503 UNAVAILABLE with MODEL_CAPACITY_EXHAUSTED - no retryDelay - use fixed wait",
|
||||
account: oauthAccount,
|
||||
body: `{
|
||||
"error": {
|
||||
@@ -635,10 +646,11 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
"message": "No capacity available for model gemini-2.5-flash on the server"
|
||||
}
|
||||
}`,
|
||||
expectedShouldRetry: false,
|
||||
expectedShouldRateLimit: true,
|
||||
minWait: 30 * time.Second,
|
||||
modelName: "gemini-2.5-flash",
|
||||
expectedShouldRetry: true,
|
||||
expectedShouldRateLimit: false,
|
||||
expectedIsModelCapacityExhausted: true,
|
||||
minWait: 1 * time.Second,
|
||||
modelName: "gemini-2.5-flash",
|
||||
},
|
||||
{
|
||||
name: "429 RESOURCE_EXHAUSTED with RATE_LIMIT_EXCEEDED - no retryDelay - use default rate limit",
|
||||
@@ -662,13 +674,16 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
shouldRetry, shouldRateLimit, wait, model := shouldTriggerAntigravitySmartRetry(tt.account, []byte(tt.body))
|
||||
shouldRetry, shouldRateLimit, wait, model, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(tt.account, []byte(tt.body))
|
||||
if shouldRetry != tt.expectedShouldRetry {
|
||||
t.Errorf("shouldRetry = %v, want %v", shouldRetry, tt.expectedShouldRetry)
|
||||
}
|
||||
if shouldRateLimit != tt.expectedShouldRateLimit {
|
||||
t.Errorf("shouldRateLimit = %v, want %v", shouldRateLimit, tt.expectedShouldRateLimit)
|
||||
}
|
||||
if isModelCapacityExhausted != tt.expectedIsModelCapacityExhausted {
|
||||
t.Errorf("isModelCapacityExhausted = %v, want %v", isModelCapacityExhausted, tt.expectedIsModelCapacityExhausted)
|
||||
}
|
||||
if shouldRetry {
|
||||
if wait < tt.minWait {
|
||||
t.Errorf("wait = %v, want >= %v", wait, tt.minWait)
|
||||
@@ -921,6 +936,22 @@ func TestIsAntigravityAccountSwitchError(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAntigravityForwardBaseURL_DefaultDaily(t *testing.T) {
|
||||
t.Setenv(antigravityForwardBaseURLEnv, "")
|
||||
|
||||
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
|
||||
defer func() {
|
||||
antigravity.BaseURLs = oldBaseURLs
|
||||
}()
|
||||
|
||||
prodURL := "https://prod.test"
|
||||
dailyURL := "https://daily.test"
|
||||
antigravity.BaseURLs = []string{dailyURL, prodURL}
|
||||
|
||||
resolved := resolveAntigravityForwardBaseURL()
|
||||
require.Equal(t, dailyURL, resolved)
|
||||
}
|
||||
|
||||
func TestAntigravityAccountSwitchError_Error(t *testing.T) {
|
||||
err := &AntigravityAccountSwitchError{
|
||||
OriginalAccountID: 789,
|
||||
|
||||
Reference in New Issue
Block a user