e4bc35151f
覆盖 Service 层和 Handler 层的所有新增逻辑： - isSingleAccountRetry context 标记检查 - handleSmartRetry 中 503 + SingleAccountRetry 分支 - handleSingleAccountRetryInPlace 原地重试逻辑 - antigravityRetryLoop 预检查跳过限流 - sleepAntigravitySingleAccountBackoff 固定延迟退避 - 端到端集成场景验证 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
905 lines
32 KiB
Go
905 lines
32 KiB
Go
//go:build unit
|
||
|
||
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 辅助函数：构造带 SingleAccountRetry 标记的 context
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func ctxWithSingleAccountRetry() context.Context {
|
||
return context.WithValue(context.Background(), ctxkey.SingleAccountRetry, true)
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 1. isSingleAccountRetry 测试
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestIsSingleAccountRetry_True(t *testing.T) {
|
||
ctx := context.WithValue(context.Background(), ctxkey.SingleAccountRetry, true)
|
||
require.True(t, isSingleAccountRetry(ctx))
|
||
}
|
||
|
||
func TestIsSingleAccountRetry_False_NoValue(t *testing.T) {
|
||
require.False(t, isSingleAccountRetry(context.Background()))
|
||
}
|
||
|
||
func TestIsSingleAccountRetry_False_ExplicitFalse(t *testing.T) {
|
||
ctx := context.WithValue(context.Background(), ctxkey.SingleAccountRetry, false)
|
||
require.False(t, isSingleAccountRetry(ctx))
|
||
}
|
||
|
||
func TestIsSingleAccountRetry_False_WrongType(t *testing.T) {
|
||
ctx := context.WithValue(context.Background(), ctxkey.SingleAccountRetry, "true")
|
||
require.False(t, isSingleAccountRetry(ctx))
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 2. 常量验证
|
||
// ---------------------------------------------------------------------------
|
||
|
||
func TestSingleAccountRetryConstants(t *testing.T) {
|
||
require.Equal(t, 3, antigravitySingleAccountSmartRetryMaxAttempts,
|
||
"单账号原地重试最多 3 次")
|
||
require.Equal(t, 15*time.Second, antigravitySingleAccountSmartRetryMaxWait,
|
||
"单次最大等待 15s")
|
||
require.Equal(t, 30*time.Second, antigravitySingleAccountSmartRetryTotalMaxWait,
|
||
"总累计等待不超过 30s")
|
||
require.Equal(t, 30*time.Second, antigravitySingleAccountMaxWait,
|
||
"预检查最大等待 30s")
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 3. handleSmartRetry + 503 + SingleAccountRetry → 走 handleSingleAccountRetryInPlace
|
||
// （而非设模型限流 + 切换账号）
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// TestHandleSmartRetry_503_LongDelay_SingleAccountRetry_RetryInPlace
|
||
// 核心场景：503 + retryDelay >= 7s + SingleAccountRetry 标记
|
||
// → 不设模型限流、不切换账号，改为原地重试
|
||
func TestHandleSmartRetry_503_LongDelay_SingleAccountRetry_RetryInPlace(t *testing.T) {
|
||
// 原地重试成功
|
||
successResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{successResp},
|
||
errors: []error{nil},
|
||
}
|
||
|
||
repo := &stubAntigravityAccountRepo{}
|
||
account := &Account{
|
||
ID: 1,
|
||
Name: "acc-single",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
// 503 + 39s >= 7s 阈值 + MODEL_CAPACITY_EXHAUSTED
|
||
respBody := []byte(`{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro-high"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "39s"}
|
||
],
|
||
"message": "No capacity available for model gemini-3-pro-high on the server"
|
||
}
|
||
}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(), // 关键：设置单账号标记
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
accountRepo: repo,
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
}
|
||
|
||
availableURLs := []string{"https://ag-1.test"}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
// 关键断言：返回 resp（原地重试成功），而非 switchError（切换账号）
|
||
require.NotNil(t, result.resp, "should return successful response from in-place retry")
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
require.Nil(t, result.switchError, "should NOT return switchError in single account mode")
|
||
require.Nil(t, result.err)
|
||
|
||
// 验证未设模型限流（单账号模式不应设限流）
|
||
require.Len(t, repo.modelRateLimitCalls, 0,
|
||
"should NOT set model rate limit in single account retry mode")
|
||
|
||
// 验证确实调用了 upstream（原地重试）
|
||
require.GreaterOrEqual(t, len(upstream.calls), 1, "should have made at least one retry call")
|
||
}
|
||
|
||
// TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches
|
||
// 对照组：503 + retryDelay >= 7s + 无 SingleAccountRetry 标记
|
||
// → 照常设模型限流 + 切换账号
|
||
func TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches(t *testing.T) {
|
||
repo := &stubAntigravityAccountRepo{}
|
||
account := &Account{
|
||
ID: 2,
|
||
Name: "acc-multi",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
}
|
||
|
||
// 503 + 39s >= 7s 阈值
|
||
respBody := []byte(`{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro-high"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "39s"}
|
||
]
|
||
}
|
||
}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: context.Background(), // 关键：无单账号标记
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
accountRepo: repo,
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
}
|
||
|
||
availableURLs := []string{"https://ag-1.test"}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
// 对照：多账号模式返回 switchError
|
||
require.NotNil(t, result.switchError, "multi-account mode should return switchError for 503")
|
||
require.Nil(t, result.resp, "should not return resp when switchError is set")
|
||
|
||
// 对照：多账号模式应设模型限流
|
||
require.Len(t, repo.modelRateLimitCalls, 1,
|
||
"multi-account mode SHOULD set model rate limit")
|
||
}
|
||
|
||
// TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches
|
||
// 边界情况：429（非 503）+ SingleAccountRetry 标记
|
||
// → 单账号原地重试仅针对 503，429 依然走切换账号逻辑
|
||
func TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches(t *testing.T) {
|
||
repo := &stubAntigravityAccountRepo{}
|
||
account := &Account{
|
||
ID: 3,
|
||
Name: "acc-429",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
}
|
||
|
||
// 429 + 15s >= 7s 阈值
|
||
respBody := []byte(`{
|
||
"error": {
|
||
"status": "RESOURCE_EXHAUSTED",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
|
||
]
|
||
}
|
||
}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusTooManyRequests, // 429，不是 503
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(), // 有单账号标记
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
accountRepo: repo,
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
}
|
||
|
||
availableURLs := []string{"https://ag-1.test"}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
// 429 即使有单账号标记，也应走切换账号
|
||
require.NotNil(t, result.switchError, "429 should still return switchError even with SingleAccountRetry")
|
||
require.Len(t, repo.modelRateLimitCalls, 1,
|
||
"429 should still set model rate limit even with SingleAccountRetry")
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 4. handleSmartRetry + 503 + 短延迟 + SingleAccountRetry → 智能重试耗尽后不设限流
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit
|
||
// 503 + retryDelay < 7s + SingleAccountRetry → 智能重试耗尽后直接返回 503，不设限流
|
||
func TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit(t *testing.T) {
|
||
// 智能重试也返回 503
|
||
failRespBody := `{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||
]
|
||
}
|
||
}`
|
||
failResp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{failResp},
|
||
errors: []error{nil},
|
||
}
|
||
|
||
repo := &stubAntigravityAccountRepo{}
|
||
account := &Account{
|
||
ID: 4,
|
||
Name: "acc-short-503",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
}
|
||
|
||
// 0.1s < 7s 阈值
|
||
respBody := []byte(`{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||
]
|
||
}
|
||
}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
accountRepo: repo,
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
}
|
||
|
||
availableURLs := []string{"https://ag-1.test"}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
// 关键断言：单账号 503 模式下，智能重试耗尽后直接返回 503 响应，不切换
|
||
require.NotNil(t, result.resp, "should return 503 response directly for single account mode")
|
||
require.Equal(t, http.StatusServiceUnavailable, result.resp.StatusCode)
|
||
require.Nil(t, result.switchError, "should NOT switch account in single account mode")
|
||
|
||
// 关键断言：不设模型限流
|
||
require.Len(t, repo.modelRateLimitCalls, 0,
|
||
"should NOT set model rate limit for 503 in single account mode")
|
||
}
|
||
|
||
// TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit
|
||
// 对照组：503 + retryDelay < 7s + 无 SingleAccountRetry → 智能重试耗尽后照常设限流
|
||
func TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit(t *testing.T) {
|
||
failRespBody := `{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||
]
|
||
}
|
||
}`
|
||
failResp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{failResp},
|
||
errors: []error{nil},
|
||
}
|
||
|
||
repo := &stubAntigravityAccountRepo{}
|
||
account := &Account{
|
||
ID: 5,
|
||
Name: "acc-multi-503",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
}
|
||
|
||
respBody := []byte(`{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||
]
|
||
}
|
||
}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: context.Background(), // 无单账号标记
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
accountRepo: repo,
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
}
|
||
|
||
availableURLs := []string{"https://ag-1.test"}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
// 对照：多账号模式应返回 switchError
|
||
require.NotNil(t, result.switchError, "multi-account mode should return switchError for 503")
|
||
// 对照：多账号模式应设模型限流
|
||
require.Len(t, repo.modelRateLimitCalls, 1,
|
||
"multi-account mode should set model rate limit")
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 5. handleSingleAccountRetryInPlace 直接测试
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// TestHandleSingleAccountRetryInPlace_Success 原地重试成功
|
||
func TestHandleSingleAccountRetryInPlace_Success(t *testing.T) {
|
||
successResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{successResp},
|
||
errors: []error{nil},
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 10,
|
||
Name: "acc-inplace-ok",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 1*time.Second, "gemini-3-pro")
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
require.NotNil(t, result.resp, "should return successful response")
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
require.Nil(t, result.switchError, "should not switch account on success")
|
||
require.Nil(t, result.err)
|
||
}
|
||
|
||
// TestHandleSingleAccountRetryInPlace_AllRetriesFail 所有重试都失败，返回 503（不设限流）
|
||
func TestHandleSingleAccountRetryInPlace_AllRetriesFail(t *testing.T) {
|
||
// 构造 3 个 503 响应（对应 3 次原地重试）
|
||
var responses []*http.Response
|
||
var errors []error
|
||
for i := 0; i < antigravitySingleAccountSmartRetryMaxAttempts; i++ {
|
||
responses = append(responses, &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||
]
|
||
}
|
||
}`)),
|
||
})
|
||
errors = append(errors, nil)
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: responses,
|
||
errors: errors,
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 11,
|
||
Name: "acc-inplace-fail",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
origBody := []byte(`{"error":{"code":503,"status":"UNAVAILABLE"}}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{"X-Test": {"original"}},
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSingleAccountRetryInPlace(params, resp, origBody, "https://ag-1.test", 1*time.Second, "gemini-3-pro")
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
// 关键：返回 503 resp，不返回 switchError
|
||
require.NotNil(t, result.resp, "should return 503 response directly")
|
||
require.Equal(t, http.StatusServiceUnavailable, result.resp.StatusCode)
|
||
require.Nil(t, result.switchError, "should NOT return switchError - let Handler handle it")
|
||
require.Nil(t, result.err)
|
||
|
||
// 验证确实重试了指定次数
|
||
require.Len(t, upstream.calls, antigravitySingleAccountSmartRetryMaxAttempts,
|
||
"should have made exactly maxAttempts retry calls")
|
||
}
|
||
|
||
// TestHandleSingleAccountRetryInPlace_WaitDurationClamped 等待时间被限制在 [min, max] 范围
|
||
func TestHandleSingleAccountRetryInPlace_WaitDurationClamped(t *testing.T) {
|
||
// 用短延迟的成功响应，只验证不 panic
|
||
successResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{successResp},
|
||
errors: []error{nil},
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 12,
|
||
Name: "acc-clamp",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
|
||
// 等待时间过大应被 clamp 到 antigravitySingleAccountSmartRetryMaxWait
|
||
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 999*time.Second, "gemini-3-pro")
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
require.NotNil(t, result.resp)
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
}
|
||
|
||
// TestHandleSingleAccountRetryInPlace_ContextCanceled context 取消时立即返回
|
||
func TestHandleSingleAccountRetryInPlace_ContextCanceled(t *testing.T) {
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{nil},
|
||
errors: []error{nil},
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 13,
|
||
Name: "acc-cancel",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
ctx = context.WithValue(ctx, ctxkey.SingleAccountRetry, true)
|
||
cancel() // 立即取消
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctx,
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 1*time.Second, "gemini-3-pro")
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
require.Error(t, result.err, "should return context error")
|
||
// 不应调用 upstream（因为在等待阶段就被取消了）
|
||
require.Len(t, upstream.calls, 0, "should not call upstream when context is canceled")
|
||
}
|
||
|
||
// TestHandleSingleAccountRetryInPlace_NetworkError_ContinuesRetry 网络错误时继续重试
|
||
func TestHandleSingleAccountRetryInPlace_NetworkError_ContinuesRetry(t *testing.T) {
|
||
successResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||
}
|
||
upstream := &mockSmartRetryUpstream{
|
||
// 第1次网络错误（nil resp），第2次成功
|
||
responses: []*http.Response{nil, successResp},
|
||
errors: []error{nil, nil},
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 14,
|
||
Name: "acc-net-retry",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 1*time.Second, "gemini-3-pro")
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
require.NotNil(t, result.resp, "should return successful response after network error recovery")
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
require.Len(t, upstream.calls, 2, "first call fails (network error), second succeeds")
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 6. antigravityRetryLoop 预检查：单账号模式跳过限流
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// TestAntigravityRetryLoop_PreCheck_SingleAccountRetry_SkipsRateLimit
|
||
// 预检查中，如果有 SingleAccountRetry 标记，即使账号已限流也跳过直接发请求
|
||
func TestAntigravityRetryLoop_PreCheck_SingleAccountRetry_SkipsRateLimit(t *testing.T) {
|
||
// 创建一个已设模型限流的账号
|
||
upstream := &recordingOKUpstream{}
|
||
account := &Account{
|
||
ID: 20,
|
||
Name: "acc-rate-limited",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Schedulable: true,
|
||
Status: StatusActive,
|
||
Concurrency: 1,
|
||
Extra: map[string]any{
|
||
modelRateLimitsKey: map[string]any{
|
||
"claude-sonnet-4-5": map[string]any{
|
||
"rate_limit_reset_at": time.Now().Add(30 * time.Second).Format(time.RFC3339),
|
||
},
|
||
},
|
||
},
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
requestedModel: "claude-sonnet-4-5",
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
})
|
||
|
||
require.NoError(t, err, "should not return error")
|
||
require.NotNil(t, result, "should return result")
|
||
require.NotNil(t, result.resp, "should have response")
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
// 关键：尽管限流了，有 SingleAccountRetry 标记时仍然到达了 upstream
|
||
require.Equal(t, 1, upstream.calls, "should have reached upstream despite rate limit")
|
||
}
|
||
|
||
// TestAntigravityRetryLoop_PreCheck_NoSingleAccountRetry_SwitchesOnRateLimit
|
||
// 对照组：无 SingleAccountRetry + 已限流 → 预检查返回 switchError
|
||
func TestAntigravityRetryLoop_PreCheck_NoSingleAccountRetry_SwitchesOnRateLimit(t *testing.T) {
|
||
upstream := &recordingOKUpstream{}
|
||
account := &Account{
|
||
ID: 21,
|
||
Name: "acc-rate-limited-multi",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Schedulable: true,
|
||
Status: StatusActive,
|
||
Concurrency: 1,
|
||
Extra: map[string]any{
|
||
modelRateLimitsKey: map[string]any{
|
||
"claude-sonnet-4-5": map[string]any{
|
||
"rate_limit_reset_at": time.Now().Add(30 * time.Second).Format(time.RFC3339),
|
||
},
|
||
},
|
||
},
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
|
||
ctx: context.Background(), // 无单账号标记
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
requestedModel: "claude-sonnet-4-5",
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
})
|
||
|
||
require.Nil(t, result, "should not return result on rate limit switch")
|
||
require.NotNil(t, err, "should return error")
|
||
|
||
var switchErr *AntigravityAccountSwitchError
|
||
require.ErrorAs(t, err, &switchErr, "should return AntigravityAccountSwitchError")
|
||
require.Equal(t, account.ID, switchErr.OriginalAccountID)
|
||
require.Equal(t, "claude-sonnet-4-5", switchErr.RateLimitedModel)
|
||
|
||
// upstream 不应被调用（预检查就短路了）
|
||
require.Equal(t, 0, upstream.calls, "upstream should NOT be called when pre-check blocks")
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 7. 端到端集成场景测试
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// TestHandleSmartRetry_503_SingleAccount_RetryInPlace_ThenSuccess_E2E
|
||
// 端到端场景：503 + 单账号 + 原地重试第2次成功
|
||
func TestHandleSmartRetry_503_SingleAccount_RetryInPlace_ThenSuccess_E2E(t *testing.T) {
|
||
// 第1次原地重试仍返回 503，第2次成功
|
||
fail503Body := `{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||
]
|
||
}
|
||
}`
|
||
resp503 := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(fail503Body)),
|
||
}
|
||
successResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||
}
|
||
|
||
upstream := &mockSmartRetryUpstream{
|
||
responses: []*http.Response{resp503, successResp},
|
||
errors: []error{nil, nil},
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 30,
|
||
Name: "acc-e2e",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
params := antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 1*time.Second, "gemini-3-pro")
|
||
|
||
require.NotNil(t, result)
|
||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||
require.NotNil(t, result.resp, "should return successful response after 2nd attempt")
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
require.Nil(t, result.switchError)
|
||
require.Len(t, upstream.calls, 2, "first 503, second OK")
|
||
}
|
||
|
||
// TestAntigravityRetryLoop_503_SingleAccount_InPlaceRetryUsed_E2E
|
||
// 通过 antigravityRetryLoop → handleSmartRetry → handleSingleAccountRetryInPlace 完整链路
|
||
func TestAntigravityRetryLoop_503_SingleAccount_InPlaceRetryUsed_E2E(t *testing.T) {
|
||
// 初始请求返回 503 + 长延迟
|
||
initial503Body := []byte(`{
|
||
"error": {
|
||
"code": 503,
|
||
"status": "UNAVAILABLE",
|
||
"details": [
|
||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "10s"}
|
||
],
|
||
"message": "No capacity available"
|
||
}
|
||
}`)
|
||
initial503Resp := &http.Response{
|
||
StatusCode: http.StatusServiceUnavailable,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(bytes.NewReader(initial503Body)),
|
||
}
|
||
|
||
// 原地重试成功
|
||
successResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{},
|
||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||
}
|
||
|
||
upstream := &mockSmartRetryUpstream{
|
||
// 第1次调用（retryLoop 主循环）返回 503
|
||
// 第2次调用（handleSingleAccountRetryInPlace 原地重试）返回 200
|
||
responses: []*http.Response{initial503Resp, successResp},
|
||
errors: []error{nil, nil},
|
||
}
|
||
|
||
repo := &stubAntigravityAccountRepo{}
|
||
account := &Account{
|
||
ID: 31,
|
||
Name: "acc-e2e-loop",
|
||
Type: AccountTypeOAuth,
|
||
Platform: PlatformAntigravity,
|
||
Schedulable: true,
|
||
Status: StatusActive,
|
||
Concurrency: 1,
|
||
}
|
||
|
||
svc := &AntigravityGatewayService{}
|
||
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
|
||
ctx: ctxWithSingleAccountRetry(),
|
||
prefix: "[test]",
|
||
account: account,
|
||
accessToken: "token",
|
||
action: "generateContent",
|
||
body: []byte(`{"input":"test"}`),
|
||
httpUpstream: upstream,
|
||
accountRepo: repo,
|
||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||
return nil
|
||
},
|
||
})
|
||
|
||
require.NoError(t, err, "should not return error on successful retry")
|
||
require.NotNil(t, result, "should return result")
|
||
require.NotNil(t, result.resp, "should return response")
|
||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||
|
||
// 验证未设模型限流
|
||
require.Len(t, repo.modelRateLimitCalls, 0,
|
||
"should NOT set model rate limit in single account retry mode")
|
||
}
|