failover_loop_test.go 42 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065
  1. package handler
  2. import (
  3. "context"
  4. "net/http"
  5. "net/http/httptest"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/gin-gonic/gin"
  10. "github.com/stretchr/testify/require"
  11. )
  12. // ---------------------------------------------------------------------------
  13. // Mock
  14. // ---------------------------------------------------------------------------
  15. // mockTempUnscheduler 记录 TempUnscheduleRetryableError 的调用信息。
  16. type mockTempUnscheduler struct {
  17. calls []tempUnscheduleCall
  18. }
  19. type tempUnscheduleCall struct {
  20. accountID int64
  21. failoverErr *service.UpstreamFailoverError
  22. }
  23. func (m *mockTempUnscheduler) TempUnscheduleRetryableError(_ context.Context, accountID int64, failoverErr *service.UpstreamFailoverError) {
  24. m.calls = append(m.calls, tempUnscheduleCall{accountID: accountID, failoverErr: failoverErr})
  25. }
  26. func TestSameAccountRetryDelayFor(t *testing.T) {
  27. capacityErr := &service.UpstreamFailoverError{RequestScopedTransient: true}
  28. for _, tc := range []struct {
  29. name string
  30. retryCount int
  31. want time.Duration
  32. }{
  33. {name: "first retry", retryCount: 1, want: 500 * time.Millisecond},
  34. {name: "second retry", retryCount: 2, want: time.Second},
  35. {name: "third retry", retryCount: 3, want: 2 * time.Second},
  36. {name: "fourth retry", retryCount: 4, want: 4 * time.Second},
  37. {name: "fifth retry", retryCount: 5, want: 8 * time.Second},
  38. {name: "capped retry", retryCount: 10, want: 8 * time.Second},
  39. } {
  40. t.Run(tc.name, func(t *testing.T) {
  41. require.Equal(t, tc.want, sameAccountRetryDelayFor(capacityErr, tc.retryCount))
  42. })
  43. }
  44. t.Run("non request scoped errors keep fixed delay", func(t *testing.T) {
  45. require.Equal(t, 500*time.Millisecond, sameAccountRetryDelayFor(&service.UpstreamFailoverError{}, 10))
  46. })
  47. t.Run("nil error keeps fixed delay", func(t *testing.T) {
  48. require.Equal(t, 500*time.Millisecond, sameAccountRetryDelayFor(nil, 10))
  49. })
  50. t.Run("explicit oauth delay wins", func(t *testing.T) {
  51. err := &service.UpstreamFailoverError{SameAccountRetryDelay: 3 * time.Second}
  52. require.Equal(t, 3*time.Second, sameAccountRetryDelayFor(err, 1))
  53. })
  54. }
  55. func TestSameAccountRetryAllowedRequiresDeadlineAndCountBudget(t *testing.T) {
  56. err := &service.UpstreamFailoverError{
  57. RetryableOnSameAccount: true,
  58. SameAccountRetryDeadline: time.Now().Add(time.Minute),
  59. }
  60. require.False(t, sameAccountRetryAllowed(err, 100, 0))
  61. require.False(t, sameAccountRetryAllowed(err, 100, maxSameAccountRetries))
  62. require.True(t, sameAccountRetryAllowed(err, 0, maxSameAccountRetries))
  63. err.SameAccountRetryDeadline = time.Now().Add(-time.Second)
  64. require.False(t, sameAccountRetryAllowed(err, 0, 100))
  65. }
  66. func TestOAuth429RetryWaitIsBoundedByBudgetNotShortened(t *testing.T) {
  67. err := &service.UpstreamFailoverError{StatusCode: 429, RetryableOnSameAccount: true, SameAccountRetryDeadline: time.Now().Add(2 * time.Minute), SameAccountRetryMax: 3, SameAccountRetryDelay: 90 * time.Second}
  68. require.True(t, sameAccountRetryAllowed(err, 0, 3))
  69. for i := 0; i < 30; i++ {
  70. delay := sameAccountRetryDelayFor(err, 1)
  71. require.GreaterOrEqual(t, delay, 90*time.Second)
  72. require.LessOrEqual(t, delay, 108*time.Second)
  73. }
  74. err.SameAccountRetryDeadline = time.Now().Add(20 * time.Second)
  75. require.False(t, sameAccountRetryAllowed(err, 0, 3), "do not squeeze a 90-second server wait into a shorter budget")
  76. err.SameAccountRetryDelay = 0
  77. err.SameAccountRetryDeadline = time.Now().Add(time.Minute)
  78. for attempt := 1; attempt <= 3; attempt++ {
  79. delay := sameAccountRetryDelayFor(err, attempt)
  80. base := time.Second * time.Duration(1<<(attempt-1))
  81. require.GreaterOrEqual(t, delay, base)
  82. require.LessOrEqual(t, delay, base+base/5)
  83. }
  84. require.False(t, sameAccountRetryAllowed(err, 3, 10))
  85. require.False(t, sameAccountRetryAllowed(err, 0, 0))
  86. }
  87. func TestSameAccountRetryAllowedRequiresOptInAndDefaultsToCountLimit(t *testing.T) {
  88. err := &service.UpstreamFailoverError{SameAccountRetryDeadline: time.Now().Add(time.Minute)}
  89. require.False(t, sameAccountRetryAllowed(err, 0, maxSameAccountRetries))
  90. err.RetryableOnSameAccount = true
  91. err.SameAccountRetryDeadline = time.Time{}
  92. require.True(t, sameAccountRetryAllowed(err, maxSameAccountRetries-1, maxSameAccountRetries))
  93. require.False(t, sameAccountRetryAllowed(err, maxSameAccountRetries, maxSameAccountRetries))
  94. }
  95. func TestSameAccountRetryAllowedHonorsErrorMaxBeforeDeadline(t *testing.T) {
  96. err := &service.UpstreamFailoverError{
  97. RetryableOnSameAccount: true,
  98. SameAccountRetryDeadline: time.Now().Add(time.Minute),
  99. SameAccountRetryMax: 1,
  100. }
  101. require.True(t, sameAccountRetryAllowed(err, 0, maxSameAccountRetries))
  102. require.False(t, sameAccountRetryAllowed(err, 1, maxSameAccountRetries))
  103. require.False(t, sameAccountRetryAllowed(err, 0, 0), "an explicit zero retry budget remains disabled")
  104. }
  105. func TestSameAccountRetryDeadlineAllows(t *testing.T) {
  106. require.True(t, sameAccountRetryDeadlineAllows(&service.UpstreamFailoverError{}))
  107. require.True(t, sameAccountRetryDeadlineAllows(&service.UpstreamFailoverError{
  108. SameAccountRetryDeadline: time.Now().Add(time.Second),
  109. }))
  110. require.False(t, sameAccountRetryDeadlineAllows(&service.UpstreamFailoverError{
  111. SameAccountRetryDeadline: time.Now().Add(-time.Second),
  112. }))
  113. }
  114. func TestEffectiveSameAccountRetryLimitHonorsErrorCapAndDisabledAccount(t *testing.T) {
  115. account := &service.Account{Type: service.AccountTypeAPIKey, Credentials: map[string]any{"pool_mode": true, "pool_mode_retry_count": float64(3)}}
  116. require.Equal(t, 1, effectiveSameAccountRetryLimit(&service.UpstreamFailoverError{SameAccountRetryMax: 1}, account))
  117. account.Credentials["pool_mode_retry_count"] = float64(0)
  118. require.Equal(t, 0, effectiveSameAccountRetryLimit(&service.UpstreamFailoverError{SameAccountRetryMax: 1}, account))
  119. }
  120. // ---------------------------------------------------------------------------
  121. // Helper
  122. // ---------------------------------------------------------------------------
  123. func newTestFailoverErr(statusCode int, retryable, forceBilling bool) *service.UpstreamFailoverError {
  124. return &service.UpstreamFailoverError{
  125. StatusCode: statusCode,
  126. RetryableOnSameAccount: retryable,
  127. ForceCacheBilling: forceBilling,
  128. }
  129. }
  130. // ---------------------------------------------------------------------------
  131. // NewFailoverState 测试
  132. // ---------------------------------------------------------------------------
  133. func TestNewFailoverState(t *testing.T) {
  134. t.Run("初始化字段正确", func(t *testing.T) {
  135. fs := NewFailoverState(5, true)
  136. require.Equal(t, 5, fs.MaxSwitches)
  137. require.Equal(t, 0, fs.SwitchCount)
  138. require.NotNil(t, fs.FailedAccountIDs)
  139. require.Empty(t, fs.FailedAccountIDs)
  140. require.NotNil(t, fs.SameAccountRetryCount)
  141. require.Empty(t, fs.SameAccountRetryCount)
  142. require.Nil(t, fs.LastFailoverErr)
  143. require.False(t, fs.ForceCacheBilling)
  144. require.True(t, fs.hasBoundSession)
  145. })
  146. t.Run("无绑定会话", func(t *testing.T) {
  147. fs := NewFailoverState(3, false)
  148. require.Equal(t, 3, fs.MaxSwitches)
  149. require.False(t, fs.hasBoundSession)
  150. })
  151. t.Run("零最大切换次数", func(t *testing.T) {
  152. fs := NewFailoverState(0, false)
  153. require.Equal(t, 0, fs.MaxSwitches)
  154. })
  155. }
  156. // ---------------------------------------------------------------------------
  157. // sleepWithContext 测试
  158. // ---------------------------------------------------------------------------
  159. func TestSleepWithContext(t *testing.T) {
  160. t.Run("零时长立即返回true", func(t *testing.T) {
  161. start := time.Now()
  162. ok := sleepWithContext(context.Background(), 0)
  163. require.True(t, ok)
  164. require.Less(t, time.Since(start), 50*time.Millisecond)
  165. })
  166. t.Run("负时长立即返回true", func(t *testing.T) {
  167. start := time.Now()
  168. ok := sleepWithContext(context.Background(), -1*time.Second)
  169. require.True(t, ok)
  170. require.Less(t, time.Since(start), 50*time.Millisecond)
  171. })
  172. t.Run("正常等待后返回true", func(t *testing.T) {
  173. start := time.Now()
  174. ok := sleepWithContext(context.Background(), 50*time.Millisecond)
  175. elapsed := time.Since(start)
  176. require.True(t, ok)
  177. require.GreaterOrEqual(t, elapsed, 40*time.Millisecond)
  178. require.Less(t, elapsed, 500*time.Millisecond)
  179. })
  180. t.Run("已取消context立即返回false", func(t *testing.T) {
  181. ctx, cancel := context.WithCancel(context.Background())
  182. cancel()
  183. start := time.Now()
  184. ok := sleepWithContext(ctx, 5*time.Second)
  185. require.False(t, ok)
  186. require.Less(t, time.Since(start), 50*time.Millisecond)
  187. })
  188. t.Run("等待期间context取消返回false", func(t *testing.T) {
  189. ctx, cancel := context.WithCancel(context.Background())
  190. go func() {
  191. time.Sleep(30 * time.Millisecond)
  192. cancel()
  193. }()
  194. start := time.Now()
  195. ok := sleepWithContext(ctx, 5*time.Second)
  196. elapsed := time.Since(start)
  197. require.False(t, ok)
  198. require.Less(t, elapsed, 500*time.Millisecond)
  199. })
  200. }
  201. // ---------------------------------------------------------------------------
  202. // HandleFailoverError — 基本切换流程
  203. // ---------------------------------------------------------------------------
  204. func TestHandleFailoverError_BasicSwitch(t *testing.T) {
  205. t.Run("显式停止不切换账号且旧错误默认仍切换", func(t *testing.T) {
  206. mock := &mockTempUnscheduler{}
  207. fs := NewFailoverState(3, false)
  208. stopErr := &service.UpstreamFailoverError{
  209. Stage: service.GatewayFailureStageAccountAuth,
  210. Scope: service.GatewayFailureScopeProvider,
  211. NextAccountAction: service.NextAccountStop,
  212. }
  213. action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformGrok, maxSameAccountRetries, stopErr)
  214. require.Equal(t, FailoverExhausted, action)
  215. require.Zero(t, fs.SwitchCount)
  216. require.Empty(t, fs.FailedAccountIDs)
  217. require.Equal(t, stopErr, fs.LastFailoverErr)
  218. legacyErr := newTestFailoverErr(http.StatusTooManyRequests, false, false)
  219. action = fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformGrok, maxSameAccountRetries, legacyErr)
  220. require.Equal(t, FailoverContinue, action)
  221. require.Equal(t, 1, fs.SwitchCount)
  222. require.Contains(t, fs.FailedAccountIDs, int64(100))
  223. })
  224. t.Run("已取消的认证失败不改变切换状态", func(t *testing.T) {
  225. ctx, cancel := context.WithCancel(context.Background())
  226. cancel()
  227. mock := &mockTempUnscheduler{}
  228. fs := NewFailoverState(3, false)
  229. err := &service.UpstreamFailoverError{
  230. Stage: service.GatewayFailureStageAccountAuth,
  231. Scope: service.GatewayFailureScopeAccount,
  232. NextAccountAction: service.NextAccountRetry,
  233. }
  234. action := fs.HandleFailoverError(ctx, mock, 101, service.PlatformGrok, maxSameAccountRetries, err)
  235. require.Equal(t, FailoverCanceled, action)
  236. require.Zero(t, fs.SwitchCount)
  237. require.Empty(t, fs.FailedAccountIDs)
  238. require.Nil(t, fs.LastFailoverErr)
  239. require.Empty(t, mock.calls)
  240. })
  241. t.Run("非重试错误_非Antigravity_直接切换", func(t *testing.T) {
  242. mock := &mockTempUnscheduler{}
  243. fs := NewFailoverState(3, false)
  244. err := newTestFailoverErr(500, false, false)
  245. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  246. require.Equal(t, FailoverContinue, action)
  247. require.Equal(t, 1, fs.SwitchCount)
  248. require.Contains(t, fs.FailedAccountIDs, int64(100))
  249. require.Equal(t, err, fs.LastFailoverErr)
  250. require.False(t, fs.ForceCacheBilling)
  251. require.Empty(t, mock.calls, "不应调用 TempUnschedule")
  252. })
  253. t.Run("非重试错误_Antigravity_第一次切换无延迟", func(t *testing.T) {
  254. // switchCount 从 0→1 时,sleepFailoverDelay(ctx, 1) 的延时 = (1-1)*1s = 0
  255. mock := &mockTempUnscheduler{}
  256. fs := NewFailoverState(3, false)
  257. err := newTestFailoverErr(500, false, false)
  258. start := time.Now()
  259. action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
  260. elapsed := time.Since(start)
  261. require.Equal(t, FailoverContinue, action)
  262. require.Equal(t, 1, fs.SwitchCount)
  263. require.Less(t, elapsed, 200*time.Millisecond, "第一次切换延迟应为 0")
  264. })
  265. t.Run("非重试错误_Antigravity_第二次切换有1秒延迟", func(t *testing.T) {
  266. // switchCount 从 1→2 时,sleepFailoverDelay(ctx, 2) 的延时 = (2-1)*1s = 1s
  267. mock := &mockTempUnscheduler{}
  268. fs := NewFailoverState(3, false)
  269. fs.SwitchCount = 1 // 模拟已切换一次
  270. err := newTestFailoverErr(500, false, false)
  271. start := time.Now()
  272. action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
  273. elapsed := time.Since(start)
  274. require.Equal(t, FailoverContinue, action)
  275. require.Equal(t, 2, fs.SwitchCount)
  276. require.GreaterOrEqual(t, elapsed, 800*time.Millisecond, "第二次切换延迟应约 1s")
  277. require.Less(t, elapsed, 3*time.Second)
  278. })
  279. t.Run("连续切换直到耗尽", func(t *testing.T) {
  280. mock := &mockTempUnscheduler{}
  281. fs := NewFailoverState(2, false)
  282. // 第一次切换:0→1
  283. err1 := newTestFailoverErr(500, false, false)
  284. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
  285. require.Equal(t, FailoverContinue, action)
  286. require.Equal(t, 1, fs.SwitchCount)
  287. // 第二次切换:1→2
  288. err2 := newTestFailoverErr(502, false, false)
  289. action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
  290. require.Equal(t, FailoverContinue, action)
  291. require.Equal(t, 2, fs.SwitchCount)
  292. // 第三次已耗尽:SwitchCount(2) >= MaxSwitches(2)
  293. err3 := newTestFailoverErr(503, false, false)
  294. action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
  295. require.Equal(t, FailoverExhausted, action)
  296. require.Equal(t, 2, fs.SwitchCount, "耗尽时不应继续递增")
  297. // 验证失败账号列表
  298. require.Len(t, fs.FailedAccountIDs, 3)
  299. require.Contains(t, fs.FailedAccountIDs, int64(100))
  300. require.Contains(t, fs.FailedAccountIDs, int64(200))
  301. require.Contains(t, fs.FailedAccountIDs, int64(300))
  302. // LastFailoverErr 应为最后一次的错误
  303. require.Equal(t, err3, fs.LastFailoverErr)
  304. })
  305. t.Run("MaxSwitches为0时首次即耗尽", func(t *testing.T) {
  306. mock := &mockTempUnscheduler{}
  307. fs := NewFailoverState(0, false)
  308. err := newTestFailoverErr(500, false, false)
  309. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  310. require.Equal(t, FailoverExhausted, action)
  311. require.Equal(t, 0, fs.SwitchCount)
  312. require.Contains(t, fs.FailedAccountIDs, int64(100))
  313. })
  314. }
  315. // ---------------------------------------------------------------------------
  316. // HandleFailoverError — 缓存计费 (ForceCacheBilling)
  317. // ---------------------------------------------------------------------------
  318. func TestHandleFailoverError_CacheBilling(t *testing.T) {
  319. t.Run("hasBoundSession为true且实际切换时设置ForceCacheBilling", func(t *testing.T) {
  320. mock := &mockTempUnscheduler{}
  321. fs := NewFailoverState(3, true) // hasBoundSession=true
  322. err := newTestFailoverErr(500, false, false)
  323. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  324. require.True(t, fs.ForceCacheBilling)
  325. })
  326. t.Run("同账号重试时仅凭hasBoundSession不设置ForceCacheBilling", func(t *testing.T) {
  327. mock := &mockTempUnscheduler{}
  328. fs := NewFailoverState(3, true)
  329. err := newTestFailoverErr(400, true, false)
  330. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  331. require.False(t, fs.ForceCacheBilling)
  332. require.Zero(t, fs.SwitchCount)
  333. })
  334. t.Run("OAuth deadline未到但次数耗尽时停止同账号重试", func(t *testing.T) {
  335. mock := &mockTempUnscheduler{}
  336. fs := NewFailoverState(3, true)
  337. fs.SameAccountRetryCount[100] = maxSameAccountRetries
  338. err := newTestFailoverErr(http.StatusTooManyRequests, true, false)
  339. err.SameAccountRetryDeadline = time.Now().Add(time.Minute)
  340. err.SameAccountRetryDelay = time.Nanosecond
  341. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  342. require.True(t, fs.ForceCacheBilling)
  343. require.Equal(t, 1, fs.SwitchCount)
  344. require.Equal(t, maxSameAccountRetries, fs.SameAccountRetryCount[100])
  345. require.Len(t, mock.calls, 1)
  346. })
  347. t.Run("同账号重试耗尽并实际切换时设置ForceCacheBilling", func(t *testing.T) {
  348. mock := &mockTempUnscheduler{}
  349. fs := NewFailoverState(3, true)
  350. err := newTestFailoverErr(400, true, false)
  351. for i := 0; i < maxSameAccountRetries; i++ {
  352. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  353. require.False(t, fs.ForceCacheBilling)
  354. }
  355. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  356. require.True(t, fs.ForceCacheBilling)
  357. require.Equal(t, 1, fs.SwitchCount)
  358. })
  359. t.Run("failoverErr.ForceCacheBilling为true时设置", func(t *testing.T) {
  360. mock := &mockTempUnscheduler{}
  361. fs := NewFailoverState(3, false)
  362. err := newTestFailoverErr(500, false, true) // ForceCacheBilling=true
  363. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  364. require.True(t, fs.ForceCacheBilling)
  365. })
  366. t.Run("同账号重试保留显式ForceCacheBilling", func(t *testing.T) {
  367. mock := &mockTempUnscheduler{}
  368. fs := NewFailoverState(3, true)
  369. err := newTestFailoverErr(400, true, true)
  370. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  371. require.True(t, fs.ForceCacheBilling)
  372. require.Zero(t, fs.SwitchCount)
  373. })
  374. t.Run("两者均为false时不设置", func(t *testing.T) {
  375. mock := &mockTempUnscheduler{}
  376. fs := NewFailoverState(3, false)
  377. err := newTestFailoverErr(500, false, false)
  378. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  379. require.False(t, fs.ForceCacheBilling)
  380. })
  381. t.Run("一旦设置不会被后续错误重置", func(t *testing.T) {
  382. mock := &mockTempUnscheduler{}
  383. fs := NewFailoverState(3, false)
  384. // 第一次:ForceCacheBilling=true → 设置
  385. err1 := newTestFailoverErr(500, false, true)
  386. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
  387. require.True(t, fs.ForceCacheBilling)
  388. // 第二次:ForceCacheBilling=false → 仍然保持 true
  389. err2 := newTestFailoverErr(502, false, false)
  390. fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
  391. require.True(t, fs.ForceCacheBilling, "ForceCacheBilling 一旦设置不应被重置")
  392. })
  393. }
  394. // ---------------------------------------------------------------------------
  395. // HandleFailoverError — 同账号重试 (RetryableOnSameAccount)
  396. // ---------------------------------------------------------------------------
  397. func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
  398. t.Run("第一次重试返回FailoverContinue", func(t *testing.T) {
  399. mock := &mockTempUnscheduler{}
  400. fs := NewFailoverState(3, false)
  401. err := newTestFailoverErr(400, true, false)
  402. start := time.Now()
  403. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  404. elapsed := time.Since(start)
  405. require.Equal(t, FailoverContinue, action)
  406. require.Equal(t, 1, fs.SameAccountRetryCount[100])
  407. require.Equal(t, 0, fs.SwitchCount, "同账号重试不应增加切换计数")
  408. require.NotContains(t, fs.FailedAccountIDs, int64(100), "同账号重试不应加入失败列表")
  409. require.Empty(t, mock.calls, "同账号重试期间不应调用 TempUnschedule")
  410. // 验证等待了 sameAccountRetryDelay (500ms)
  411. require.GreaterOrEqual(t, elapsed, 400*time.Millisecond)
  412. require.Less(t, elapsed, 2*time.Second)
  413. })
  414. t.Run("达到最大重试次数前均返回FailoverContinue", func(t *testing.T) {
  415. mock := &mockTempUnscheduler{}
  416. fs := NewFailoverState(3, false)
  417. err := newTestFailoverErr(400, true, false)
  418. for i := 1; i <= maxSameAccountRetries; i++ {
  419. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  420. require.Equal(t, FailoverContinue, action)
  421. require.Equal(t, i, fs.SameAccountRetryCount[100])
  422. }
  423. require.Empty(t, mock.calls, "达到最大重试次数前均不应调用 TempUnschedule")
  424. })
  425. t.Run("超过最大重试次数后触发TempUnschedule并切换", func(t *testing.T) {
  426. mock := &mockTempUnscheduler{}
  427. fs := NewFailoverState(3, false)
  428. err := newTestFailoverErr(400, true, false)
  429. for i := 0; i < maxSameAccountRetries; i++ {
  430. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  431. }
  432. require.Equal(t, maxSameAccountRetries, fs.SameAccountRetryCount[100])
  433. // 第 maxSameAccountRetries+1 次:重试耗尽,应切换账号
  434. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  435. require.Equal(t, FailoverContinue, action)
  436. require.Equal(t, 1, fs.SwitchCount)
  437. require.Contains(t, fs.FailedAccountIDs, int64(100))
  438. // 验证 TempUnschedule 被调用
  439. require.Len(t, mock.calls, 1)
  440. require.Equal(t, int64(100), mock.calls[0].accountID)
  441. require.Equal(t, err, mock.calls[0].failoverErr)
  442. })
  443. t.Run("不同账号独立跟踪重试次数", func(t *testing.T) {
  444. mock := &mockTempUnscheduler{}
  445. fs := NewFailoverState(5, false)
  446. err := newTestFailoverErr(400, true, false)
  447. // 账号 100 第一次重试
  448. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  449. require.Equal(t, FailoverContinue, action)
  450. require.Equal(t, 1, fs.SameAccountRetryCount[100])
  451. // 账号 200 第一次重试(独立计数)
  452. action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err)
  453. require.Equal(t, FailoverContinue, action)
  454. require.Equal(t, 1, fs.SameAccountRetryCount[200])
  455. require.Equal(t, 1, fs.SameAccountRetryCount[100], "账号 100 的计数不应受影响")
  456. })
  457. t.Run("重试耗尽后再次遇到同账号_直接切换", func(t *testing.T) {
  458. mock := &mockTempUnscheduler{}
  459. fs := NewFailoverState(5, false)
  460. err := newTestFailoverErr(400, true, false)
  461. // 耗尽账号 100 的重试
  462. for i := 0; i < maxSameAccountRetries; i++ {
  463. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  464. }
  465. // 第 maxSameAccountRetries+1 次: 重试耗尽 → 切换
  466. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  467. require.Equal(t, FailoverContinue, action)
  468. // 再次遇到账号 100,计数仍为 maxSameAccountRetries,条件不满足 → 直接切换
  469. action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  470. require.Equal(t, FailoverContinue, action)
  471. require.Len(t, mock.calls, 2, "第二次耗尽也应调用 TempUnschedule")
  472. })
  473. t.Run("尊重账号级retryLimit_配置1次只重试1次", func(t *testing.T) {
  474. // 回归测试:Anthropic 等路径此前硬编码同账号重试 3 次,忽略账号
  475. // pool_mode_retry_count 配置。此处验证传入 retryLimit=1 时只重试 1 次即切换。
  476. mock := &mockTempUnscheduler{}
  477. fs := NewFailoverState(5, false)
  478. err := newTestFailoverErr(403, true, false)
  479. const retryLimit = 1
  480. // 第 1 次:同账号重试
  481. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
  482. require.Equal(t, FailoverContinue, action)
  483. require.Equal(t, 1, fs.SameAccountRetryCount[100])
  484. require.Equal(t, 0, fs.SwitchCount, "首次重试不应切换账号")
  485. require.Empty(t, mock.calls, "未耗尽前不应 TempUnschedule")
  486. // 第 2 次:已达上限 1 → 不再同账号重试,直接切换 + TempUnschedule
  487. action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
  488. require.Equal(t, FailoverContinue, action)
  489. require.Equal(t, 1, fs.SameAccountRetryCount[100], "重试计数不应超过 retryLimit")
  490. require.Equal(t, 1, fs.SwitchCount, "重试耗尽应切换账号")
  491. require.Contains(t, fs.FailedAccountIDs, int64(100))
  492. require.Len(t, mock.calls, 1, "重试耗尽应触发 TempUnschedule")
  493. })
  494. t.Run("retryLimit为0时立即切换不重试", func(t *testing.T) {
  495. // pool_mode_retry_count=0 表示关闭同账号重试(如 GPT Image 账号)。
  496. mock := &mockTempUnscheduler{}
  497. fs := NewFailoverState(5, false)
  498. err := newTestFailoverErr(403, true, false)
  499. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", 0, err)
  500. require.Equal(t, FailoverContinue, action)
  501. require.Equal(t, 0, fs.SameAccountRetryCount[100], "retryLimit=0 不应发生同账号重试")
  502. require.Equal(t, 1, fs.SwitchCount, "应立即切换账号")
  503. require.Len(t, mock.calls, 1, "应立即 TempUnschedule")
  504. })
  505. }
  506. // ---------------------------------------------------------------------------
  507. // HandleFailoverError — TempUnschedule 调用验证
  508. // ---------------------------------------------------------------------------
  509. func TestHandleFailoverError_TempUnschedule(t *testing.T) {
  510. t.Run("非重试错误不调用TempUnschedule", func(t *testing.T) {
  511. mock := &mockTempUnscheduler{}
  512. fs := NewFailoverState(3, false)
  513. err := newTestFailoverErr(500, false, false) // RetryableOnSameAccount=false
  514. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  515. require.Empty(t, mock.calls)
  516. })
  517. t.Run("重试错误耗尽后调用TempUnschedule_传入正确参数", func(t *testing.T) {
  518. mock := &mockTempUnscheduler{}
  519. fs := NewFailoverState(3, false)
  520. err := newTestFailoverErr(502, true, false)
  521. for i := 0; i < maxSameAccountRetries; i++ {
  522. fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
  523. }
  524. // 再次触发时才会执行 TempUnschedule + 切换
  525. fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
  526. require.Len(t, mock.calls, 1)
  527. require.Equal(t, int64(42), mock.calls[0].accountID)
  528. require.Equal(t, 502, mock.calls[0].failoverErr.StatusCode)
  529. require.True(t, mock.calls[0].failoverErr.RetryableOnSameAccount)
  530. })
  531. }
  532. // ---------------------------------------------------------------------------
  533. // HandleFailoverError — Context 取消
  534. // ---------------------------------------------------------------------------
  535. func TestHandleFailoverError_ContextCanceled(t *testing.T) {
  536. t.Run("同账号重试sleep期间context取消", func(t *testing.T) {
  537. mock := &mockTempUnscheduler{}
  538. fs := NewFailoverState(3, false)
  539. err := newTestFailoverErr(400, true, false)
  540. ctx, cancel := context.WithCancel(context.Background())
  541. go func() {
  542. time.Sleep(30 * time.Millisecond)
  543. cancel() // 通过入口检查后、sleep 期间取消
  544. }()
  545. start := time.Now()
  546. action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
  547. elapsed := time.Since(start)
  548. require.Equal(t, FailoverCanceled, action)
  549. require.Less(t, elapsed, 400*time.Millisecond, "sleep 应被取消打断")
  550. // 进入重试分支后才取消:重试计数已递增
  551. require.Equal(t, 1, fs.SameAccountRetryCount[100])
  552. })
  553. t.Run("入口即已取消_不改动任何failover状态", func(t *testing.T) {
  554. mock := &mockTempUnscheduler{}
  555. fs := NewFailoverState(3, false)
  556. err := newTestFailoverErr(520, false, false)
  557. ctx, cancel := context.WithCancel(context.Background())
  558. cancel() // 调用前客户端已断开
  559. start := time.Now()
  560. action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
  561. elapsed := time.Since(start)
  562. require.Equal(t, FailoverCanceled, action)
  563. require.Less(t, elapsed, 100*time.Millisecond, "应立即返回")
  564. // 入口已取消时不得改变任何 failover 状态。
  565. require.Equal(t, 0, fs.SwitchCount, "取消的请求不应计入切换")
  566. require.Equal(t, 0, fs.SameAccountRetryCount[100], "取消的请求不应改动重试计数")
  567. require.NotContains(t, fs.FailedAccountIDs, int64(100))
  568. require.Nil(t, fs.LastFailoverErr)
  569. require.Empty(t, mock.calls, "不应触发 TempUnschedule")
  570. })
  571. t.Run("Antigravity延迟期间context取消", func(t *testing.T) {
  572. mock := &mockTempUnscheduler{}
  573. fs := NewFailoverState(3, false)
  574. fs.SwitchCount = 1 // 下一次 switchCount=2 → delay = 1s
  575. err := newTestFailoverErr(500, false, false)
  576. ctx, cancel := context.WithCancel(context.Background())
  577. cancel() // 立即取消
  578. start := time.Now()
  579. action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
  580. elapsed := time.Since(start)
  581. require.Equal(t, FailoverCanceled, action)
  582. require.Less(t, elapsed, 100*time.Millisecond, "应立即返回而非等待 1s")
  583. })
  584. }
  585. // ---------------------------------------------------------------------------
  586. // HandleFailoverError — FailedAccountIDs 跟踪
  587. // ---------------------------------------------------------------------------
  588. func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
  589. t.Run("切换时添加到失败列表", func(t *testing.T) {
  590. mock := &mockTempUnscheduler{}
  591. fs := NewFailoverState(3, false)
  592. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
  593. require.Contains(t, fs.FailedAccountIDs, int64(100))
  594. fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, newTestFailoverErr(502, false, false))
  595. require.Contains(t, fs.FailedAccountIDs, int64(200))
  596. require.Len(t, fs.FailedAccountIDs, 2)
  597. })
  598. t.Run("耗尽时也添加到失败列表", func(t *testing.T) {
  599. mock := &mockTempUnscheduler{}
  600. fs := NewFailoverState(0, false)
  601. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
  602. require.Equal(t, FailoverExhausted, action)
  603. require.Contains(t, fs.FailedAccountIDs, int64(100))
  604. })
  605. t.Run("同账号重试期间不添加到失败列表", func(t *testing.T) {
  606. mock := &mockTempUnscheduler{}
  607. fs := NewFailoverState(3, false)
  608. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(400, true, false))
  609. require.Equal(t, FailoverContinue, action)
  610. require.NotContains(t, fs.FailedAccountIDs, int64(100))
  611. })
  612. t.Run("同一账号多次切换不重复添加", func(t *testing.T) {
  613. mock := &mockTempUnscheduler{}
  614. fs := NewFailoverState(5, false)
  615. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
  616. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
  617. require.Len(t, fs.FailedAccountIDs, 1, "map 天然去重")
  618. })
  619. }
  620. // ---------------------------------------------------------------------------
  621. // HandleFailoverError — LastFailoverErr 更新
  622. // ---------------------------------------------------------------------------
  623. func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
  624. t.Run("每次调用都更新LastFailoverErr", func(t *testing.T) {
  625. mock := &mockTempUnscheduler{}
  626. fs := NewFailoverState(3, false)
  627. err1 := newTestFailoverErr(500, false, false)
  628. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
  629. require.Equal(t, err1, fs.LastFailoverErr)
  630. err2 := newTestFailoverErr(502, false, false)
  631. fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
  632. require.Equal(t, err2, fs.LastFailoverErr)
  633. })
  634. t.Run("同账号重试时也更新LastFailoverErr", func(t *testing.T) {
  635. mock := &mockTempUnscheduler{}
  636. fs := NewFailoverState(3, false)
  637. err := newTestFailoverErr(400, true, false)
  638. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  639. require.Equal(t, err, fs.LastFailoverErr)
  640. })
  641. }
  642. // ---------------------------------------------------------------------------
  643. // HandleFailoverError — 综合集成场景
  644. // ---------------------------------------------------------------------------
  645. func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
  646. t.Run("模拟完整failover流程_多账号混合重试与切换", func(t *testing.T) {
  647. mock := &mockTempUnscheduler{}
  648. fs := NewFailoverState(3, true) // hasBoundSession=true
  649. // 1. 账号 100 遇到可重试错误,同账号重试 maxSameAccountRetries 次
  650. retryErr := newTestFailoverErr(400, true, false)
  651. for i := 0; i < maxSameAccountRetries; i++ {
  652. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
  653. require.Equal(t, FailoverContinue, action)
  654. require.False(t, fs.ForceCacheBilling, "同账号重试期间不应仅因绑定会话强制缓存计费")
  655. }
  656. // 2. 账号 100 超过重试上限 → TempUnschedule + 切换
  657. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
  658. require.Equal(t, FailoverContinue, action)
  659. require.Equal(t, 1, fs.SwitchCount)
  660. require.True(t, fs.ForceCacheBilling, "实际切换账号时应设置 ForceCacheBilling")
  661. require.Len(t, mock.calls, 1)
  662. // 3. 账号 200 遇到不可重试错误 → 直接切换
  663. switchErr := newTestFailoverErr(500, false, false)
  664. action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, switchErr)
  665. require.Equal(t, FailoverContinue, action)
  666. require.Equal(t, 2, fs.SwitchCount)
  667. // 4. 账号 300 遇到不可重试错误 → 再切换
  668. action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, switchErr)
  669. require.Equal(t, FailoverContinue, action)
  670. require.Equal(t, 3, fs.SwitchCount)
  671. // 5. 账号 400 → 已耗尽 (SwitchCount=3 >= MaxSwitches=3)
  672. action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", maxSameAccountRetries, switchErr)
  673. require.Equal(t, FailoverExhausted, action)
  674. // 最终状态验证
  675. require.Equal(t, 3, fs.SwitchCount, "耗尽时不再递增")
  676. require.Len(t, fs.FailedAccountIDs, 4, "4个不同账号都在失败列表中")
  677. require.True(t, fs.ForceCacheBilling)
  678. require.Len(t, mock.calls, 1, "只有账号 100 触发了 TempUnschedule")
  679. })
  680. t.Run("模拟Antigravity平台完整流程", func(t *testing.T) {
  681. mock := &mockTempUnscheduler{}
  682. fs := NewFailoverState(2, false)
  683. err := newTestFailoverErr(500, false, false)
  684. // 第一次切换:delay = 0s
  685. start := time.Now()
  686. action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
  687. elapsed := time.Since(start)
  688. require.Equal(t, FailoverContinue, action)
  689. require.Less(t, elapsed, 200*time.Millisecond, "第一次切换延迟为 0")
  690. // 第二次切换:delay = 1s
  691. start = time.Now()
  692. action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
  693. elapsed = time.Since(start)
  694. require.Equal(t, FailoverContinue, action)
  695. require.GreaterOrEqual(t, elapsed, 800*time.Millisecond, "第二次切换延迟约 1s")
  696. // 第三次:耗尽(无延迟,因为在检查延迟之前就返回了)
  697. start = time.Now()
  698. action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, maxSameAccountRetries, err)
  699. elapsed = time.Since(start)
  700. require.Equal(t, FailoverExhausted, action)
  701. require.Less(t, elapsed, 200*time.Millisecond, "耗尽时不应有延迟")
  702. })
  703. t.Run("ForceCacheBilling通过错误标志设置", func(t *testing.T) {
  704. mock := &mockTempUnscheduler{}
  705. fs := NewFailoverState(3, false) // hasBoundSession=false
  706. // 第一次:ForceCacheBilling=false
  707. err1 := newTestFailoverErr(500, false, false)
  708. fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
  709. require.False(t, fs.ForceCacheBilling)
  710. // 第二次:ForceCacheBilling=true(Antigravity 粘性会话切换)
  711. err2 := newTestFailoverErr(500, false, true)
  712. fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
  713. require.True(t, fs.ForceCacheBilling, "错误标志应触发 ForceCacheBilling")
  714. // 第三次:ForceCacheBilling=false,但状态仍保持 true
  715. err3 := newTestFailoverErr(500, false, false)
  716. fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
  717. require.True(t, fs.ForceCacheBilling, "不应重置")
  718. })
  719. }
  720. // ---------------------------------------------------------------------------
  721. // HandleFailoverError — 边界条件
  722. // ---------------------------------------------------------------------------
  723. func TestHandleFailoverError_EdgeCases(t *testing.T) {
  724. t.Run("StatusCode为0的错误也能正常处理", func(t *testing.T) {
  725. mock := &mockTempUnscheduler{}
  726. fs := NewFailoverState(3, false)
  727. err := newTestFailoverErr(0, false, false)
  728. action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
  729. require.Equal(t, FailoverContinue, action)
  730. })
  731. t.Run("AccountID为0也能正常跟踪", func(t *testing.T) {
  732. mock := &mockTempUnscheduler{}
  733. fs := NewFailoverState(3, false)
  734. err := newTestFailoverErr(500, true, false)
  735. action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", maxSameAccountRetries, err)
  736. require.Equal(t, FailoverContinue, action)
  737. require.Equal(t, 1, fs.SameAccountRetryCount[0])
  738. })
  739. t.Run("负AccountID也能正常跟踪", func(t *testing.T) {
  740. mock := &mockTempUnscheduler{}
  741. fs := NewFailoverState(3, false)
  742. err := newTestFailoverErr(500, true, false)
  743. action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", maxSameAccountRetries, err)
  744. require.Equal(t, FailoverContinue, action)
  745. require.Equal(t, 1, fs.SameAccountRetryCount[-1])
  746. })
  747. t.Run("空平台名称不触发Antigravity延迟", func(t *testing.T) {
  748. mock := &mockTempUnscheduler{}
  749. fs := NewFailoverState(3, false)
  750. fs.SwitchCount = 1
  751. err := newTestFailoverErr(500, false, false)
  752. start := time.Now()
  753. action := fs.HandleFailoverError(context.Background(), mock, 100, "", maxSameAccountRetries, err)
  754. elapsed := time.Since(start)
  755. require.Equal(t, FailoverContinue, action)
  756. require.Less(t, elapsed, 200*time.Millisecond, "空平台不应触发 Antigravity 延迟")
  757. })
  758. }
  759. // ---------------------------------------------------------------------------
  760. // HandleSelectionExhausted 测试
  761. // ---------------------------------------------------------------------------
  762. func TestHandleSelectionExhausted(t *testing.T) {
  763. t.Run("无LastFailoverErr时返回Exhausted", func(t *testing.T) {
  764. fs := NewFailoverState(3, false)
  765. // LastFailoverErr 为 nil
  766. action := fs.HandleSelectionExhausted(context.Background())
  767. require.Equal(t, FailoverExhausted, action)
  768. })
  769. t.Run("非503错误返回Exhausted", func(t *testing.T) {
  770. fs := NewFailoverState(3, false)
  771. fs.LastFailoverErr = newTestFailoverErr(500, false, false)
  772. action := fs.HandleSelectionExhausted(context.Background())
  773. require.Equal(t, FailoverExhausted, action)
  774. })
  775. t.Run("503且未耗尽_等待后返回Continue并清除失败列表", func(t *testing.T) {
  776. fs := NewFailoverState(3, false)
  777. fs.LastFailoverErr = newTestFailoverErr(503, false, false)
  778. fs.FailedAccountIDs[100] = struct{}{}
  779. fs.SwitchCount = 1
  780. start := time.Now()
  781. action := fs.HandleSelectionExhausted(context.Background())
  782. elapsed := time.Since(start)
  783. require.Equal(t, FailoverContinue, action)
  784. require.Empty(t, fs.FailedAccountIDs, "应清除失败账号列表")
  785. require.GreaterOrEqual(t, elapsed, 1500*time.Millisecond, "应等待约 2s")
  786. require.Less(t, elapsed, 5*time.Second)
  787. })
  788. t.Run("503但SwitchCount已超过MaxSwitches_返回Exhausted", func(t *testing.T) {
  789. fs := NewFailoverState(2, false)
  790. fs.LastFailoverErr = newTestFailoverErr(503, false, false)
  791. fs.SwitchCount = 3 // > MaxSwitches(2)
  792. start := time.Now()
  793. action := fs.HandleSelectionExhausted(context.Background())
  794. elapsed := time.Since(start)
  795. require.Equal(t, FailoverExhausted, action)
  796. require.Less(t, elapsed, 100*time.Millisecond, "不应等待")
  797. })
  798. t.Run("503但context已取消_返回Canceled", func(t *testing.T) {
  799. fs := NewFailoverState(3, false)
  800. fs.LastFailoverErr = newTestFailoverErr(503, false, false)
  801. ctx, cancel := context.WithCancel(context.Background())
  802. cancel()
  803. start := time.Now()
  804. action := fs.HandleSelectionExhausted(ctx)
  805. elapsed := time.Since(start)
  806. require.Equal(t, FailoverCanceled, action)
  807. require.Less(t, elapsed, 100*time.Millisecond, "应立即返回")
  808. })
  809. t.Run("context已取消_非503也返回Canceled而非Exhausted", func(t *testing.T) {
  810. // #4257 核心场景:客户端断开后选号失败源于 context canceled,
  811. // 不应被当成账号耗尽转成 502。
  812. fs := NewFailoverState(3, false)
  813. fs.LastFailoverErr = newTestFailoverErr(520, false, false)
  814. ctx, cancel := context.WithCancel(context.Background())
  815. cancel()
  816. action := fs.HandleSelectionExhausted(ctx)
  817. require.Equal(t, FailoverCanceled, action)
  818. })
  819. t.Run("context已取消_无LastFailoverErr也返回Canceled", func(t *testing.T) {
  820. fs := NewFailoverState(3, false)
  821. ctx, cancel := context.WithCancel(context.Background())
  822. cancel()
  823. action := fs.HandleSelectionExhausted(ctx)
  824. require.Equal(t, FailoverCanceled, action)
  825. })
  826. t.Run("503且SwitchCount等于MaxSwitches_仍可重试", func(t *testing.T) {
  827. fs := NewFailoverState(2, false)
  828. fs.LastFailoverErr = newTestFailoverErr(503, false, false)
  829. fs.SwitchCount = 2 // == MaxSwitches,条件是 <=,仍可重试
  830. action := fs.HandleSelectionExhausted(context.Background())
  831. require.Equal(t, FailoverContinue, action)
  832. })
  833. }
  834. // ---------------------------------------------------------------------------
  835. // failoverClientGone 测试
  836. // ---------------------------------------------------------------------------
  837. func TestFailoverClientGone(t *testing.T) {
  838. gin.SetMode(gin.TestMode)
  839. t.Run("活跃请求返回false", func(t *testing.T) {
  840. rec := httptest.NewRecorder()
  841. c, _ := gin.CreateTestContext(rec)
  842. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
  843. require.False(t, failoverClientGone(c))
  844. require.Equal(t, http.StatusOK, c.Writer.Status(), "不应改动状态码")
  845. })
  846. t.Run("客户端已断开_返回true并标记499", func(t *testing.T) {
  847. rec := httptest.NewRecorder()
  848. c, _ := gin.CreateTestContext(rec)
  849. ctx, cancel := context.WithCancel(context.Background())
  850. cancel()
  851. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
  852. require.True(t, failoverClientGone(c))
  853. require.Equal(t, statusClientClosedRequest, c.Writer.Status())
  854. })
  855. t.Run("响应已提交_不改状态码", func(t *testing.T) {
  856. rec := httptest.NewRecorder()
  857. c, _ := gin.CreateTestContext(rec)
  858. ctx, cancel := context.WithCancel(context.Background())
  859. cancel()
  860. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
  861. c.String(http.StatusOK, "partial")
  862. require.True(t, failoverClientGone(c))
  863. require.Equal(t, http.StatusOK, c.Writer.Status(), "已提交的状态码不应被覆盖")
  864. })
  865. t.Run("nil安全", func(t *testing.T) {
  866. require.False(t, failoverClientGone(nil))
  867. })
  868. }