channel_monitor_quota_fetcher_test.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "sync"
  7. "testing"
  8. "time"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. "github.com/Wei-Shaw/sub2api/internal/domain"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  12. "github.com/stretchr/testify/require"
  13. )
  14. // --- fetcher 依赖 stub ---
  15. type stubMonitorUsageSource struct {
  16. usage *UsageInfo
  17. err error
  18. // block 非 nil 时 GetUsageForAccount 阻塞在该 channel 上,用于并发/singleflight 测试。
  19. block chan struct{}
  20. mu sync.Mutex
  21. calls int
  22. lastCtx context.Context
  23. lastAccount *Account
  24. }
  25. func (s *stubMonitorUsageSource) GetUsageForAccount(ctx context.Context, account *Account, force ...bool) (*UsageInfo, error) {
  26. s.mu.Lock()
  27. s.calls++
  28. s.lastCtx = ctx
  29. s.lastAccount = account
  30. s.mu.Unlock()
  31. if s.block != nil {
  32. <-s.block
  33. }
  34. return s.usage, s.err
  35. }
  36. func (s *stubMonitorUsageSource) getCalls() int {
  37. s.mu.Lock()
  38. defer s.mu.Unlock()
  39. return s.calls
  40. }
  41. func (s *stubMonitorUsageSource) getLastAccount() *Account {
  42. s.mu.Lock()
  43. defer s.mu.Unlock()
  44. return s.lastAccount
  45. }
  46. type stubMonitorCNQuotaSource struct {
  47. result *CNProviderQuotaProbeResult
  48. err error
  49. calls int
  50. lastAccount *Account
  51. }
  52. func (s *stubMonitorCNQuotaSource) QueryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
  53. s.calls++
  54. s.lastAccount = account
  55. return s.result, s.err
  56. }
  57. type stubMonitorCNBalanceSource struct {
  58. result *CNProviderBalanceResult
  59. err error
  60. calls int
  61. lastAccount *Account
  62. }
  63. func (s *stubMonitorCNBalanceSource) QueryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
  64. s.calls++
  65. s.lastAccount = account
  66. return s.result, s.err
  67. }
  68. type stubMonitorAccountSource struct {
  69. accounts map[int64]*Account
  70. err error
  71. calls int
  72. }
  73. func (s *stubMonitorAccountSource) GetByID(ctx context.Context, id int64) (*Account, error) {
  74. s.calls++
  75. if s.err != nil {
  76. return nil, s.err
  77. }
  78. return s.accounts[id], nil
  79. }
  80. func newQuotaFetcherTestSetup(t *testing.T) (*ChannelMonitorQuotaFetcher, *stubMonitorUsageSource, *stubMonitorCNQuotaSource, *stubMonitorCNBalanceSource, *stubMonitorAccountSource) {
  81. t.Helper()
  82. usage := &stubMonitorUsageSource{}
  83. cnQuota := &stubMonitorCNQuotaSource{}
  84. cnBalance := &stubMonitorCNBalanceSource{}
  85. accounts := &stubMonitorAccountSource{accounts: make(map[int64]*Account)}
  86. fetcher := &ChannelMonitorQuotaFetcher{
  87. usage: usage,
  88. cnQuota: cnQuota,
  89. cnBalance: cnBalance,
  90. accounts: accounts,
  91. balanceThreshold: monitorBalanceThreshold(nil),
  92. cache: make(map[int64]monitorQuotaCacheEntry),
  93. }
  94. return fetcher, usage, cnQuota, cnBalance, accounts
  95. }
  96. // --- 分派 ---
  97. func TestQuotaFetcher_OverseasAccountUsesUsageService(t *testing.T) {
  98. fetcher, usage, _, cnQuota, accounts := newQuotaFetcherTestSetup(t)
  99. accounts.accounts[7] = &Account{ID: 7, Platform: domain.PlatformAnthropic}
  100. resets := time.Now().Add(2 * time.Hour).UTC()
  101. usage.usage = &UsageInfo{
  102. FiveHour: &UsageProgress{Utilization: 42.5, UsedRequests: 17, LimitRequests: 40, ResetsAt: &resets},
  103. SevenDay: &UsageProgress{Utilization: 10},
  104. SubscriptionTier: "PRO",
  105. }
  106. snapshot := fetcher.Fetch(context.Background(), 7)
  107. require.True(t, snapshot.Success)
  108. require.Equal(t, "usage", snapshot.Source)
  109. require.Equal(t, "PRO", snapshot.PlanLevel)
  110. require.False(t, snapshot.CredentialInvalid)
  111. require.Empty(t, snapshot.Error)
  112. require.Len(t, snapshot.Tiers, 2)
  113. fiveHour := snapshot.Tiers[0]
  114. require.Equal(t, "5h", fiveHour.Window)
  115. require.Empty(t, fiveHour.Label)
  116. require.InDelta(t, 42.5, fiveHour.UsedPercent, 0.001)
  117. require.Equal(t, float64(17), fiveHour.Used)
  118. require.Equal(t, float64(40), fiveHour.Limit)
  119. require.NotEmpty(t, fiveHour.ResetAt)
  120. require.Equal(t, "7d", snapshot.Tiers[1].Window)
  121. require.Equal(t, 1, usage.getCalls())
  122. require.Equal(t, 0, cnQuota.calls)
  123. }
  124. func TestQuotaFetcher_CodingPlanAccountUsesCNQuota(t *testing.T) {
  125. fetcher, _, cnQuota, cnBalance, accounts := newQuotaFetcherTestSetup(t)
  126. accounts.accounts[9] = &Account{
  127. ID: 9,
  128. Platform: domain.PlatformKimi,
  129. Credentials: map[string]any{"account_mode": AccountModeCoding},
  130. }
  131. cnQuota.result = &CNProviderQuotaProbeResult{
  132. Success: true,
  133. CredentialValid: true,
  134. PlanLevel: "",
  135. Tiers: []CNQuotaTier{
  136. {Window: "5h", UsedPercent: 33.3, ResetAt: "2026-08-18T06:00:00Z"},
  137. {Window: "weekly", UsedPercent: 12},
  138. },
  139. }
  140. snapshot := fetcher.Fetch(context.Background(), 9)
  141. require.True(t, snapshot.Success)
  142. require.Equal(t, "cn_quota", snapshot.Source)
  143. require.Len(t, snapshot.Tiers, 2)
  144. require.Equal(t, "5h", snapshot.Tiers[0].Window)
  145. require.InDelta(t, 33.3, snapshot.Tiers[0].UsedPercent, 0.001)
  146. require.Equal(t, "weekly", snapshot.Tiers[1].Window)
  147. require.Equal(t, 1, cnQuota.calls)
  148. require.Equal(t, 0, cnBalance.calls)
  149. }
  150. func TestQuotaFetcher_MiniMaxCodingPlanUsesCNQuota(t *testing.T) {
  151. fetcher, usage, cnQuota, cnBalance, accounts := newQuotaFetcherTestSetup(t)
  152. accounts.accounts[19] = &Account{
  153. ID: 19,
  154. Platform: domain.PlatformMiniMax,
  155. Credentials: map[string]any{"account_mode": AccountModeCoding},
  156. }
  157. cnQuota.result = &CNProviderQuotaProbeResult{
  158. Success: true,
  159. CredentialValid: true,
  160. Tiers: []CNQuotaTier{
  161. {Window: "5h", UsedPercent: 12},
  162. },
  163. }
  164. snapshot := fetcher.Fetch(context.Background(), 19)
  165. require.True(t, snapshot.Success)
  166. require.Equal(t, "cn_quota", snapshot.Source)
  167. require.Equal(t, 1, cnQuota.calls)
  168. require.Equal(t, 0, cnBalance.calls)
  169. require.Equal(t, 0, usage.getCalls())
  170. }
  171. func TestQuotaFetcher_PayGAccountUsesCNBalance(t *testing.T) {
  172. fetcher, _, _, cnBalance, accounts := newQuotaFetcherTestSetup(t)
  173. accounts.accounts[11] = &Account{
  174. ID: 11,
  175. Platform: domain.PlatformDeepseek,
  176. Credentials: map[string]any{"account_mode": AccountModePayG},
  177. }
  178. cnBalance.result = &CNProviderBalanceResult{
  179. Success: true,
  180. Available: true,
  181. Balance: 12.34,
  182. Currency: "CNY",
  183. Balances: []CNProviderBalanceEntry{
  184. {Currency: "CNY", Balance: 12.34},
  185. {Currency: "USD", Balance: 1.5},
  186. },
  187. }
  188. snapshot := fetcher.Fetch(context.Background(), 11)
  189. require.True(t, snapshot.Success)
  190. require.Equal(t, "cn_balance", snapshot.Source)
  191. require.NotNil(t, snapshot.Balance)
  192. require.InDelta(t, 12.34, *snapshot.Balance, 0.001)
  193. require.Equal(t, "CNY", snapshot.Currency)
  194. require.Len(t, snapshot.Balances, 2)
  195. require.Equal(t, "USD", snapshot.Balances[1].Currency)
  196. require.False(t, snapshot.BalanceLow)
  197. require.Empty(t, snapshot.Error)
  198. }
  199. // P2-6:fetchUncached 只 GetByID 一次,已加载的 account 指针直传数据源,
  200. // 三条路由都不能让下游重载账号。
  201. func TestQuotaFetcher_LoadsAccountOnceAndPassesItThrough(t *testing.T) {
  202. t.Run("overseas usage", func(t *testing.T) {
  203. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  204. acc := &Account{ID: 21, Platform: domain.PlatformAnthropic}
  205. accounts.accounts[21] = acc
  206. usage.usage = &UsageInfo{}
  207. fetcher.Fetch(context.Background(), 21)
  208. require.Equal(t, 1, accounts.calls)
  209. require.Same(t, acc, usage.getLastAccount())
  210. require.Equal(t, 1, usage.getCalls())
  211. })
  212. t.Run("cn coding plan", func(t *testing.T) {
  213. fetcher, _, cnQuota, _, accounts := newQuotaFetcherTestSetup(t)
  214. acc := &Account{ID: 22, Platform: domain.PlatformKimi, Credentials: map[string]any{"account_mode": AccountModeCoding}}
  215. accounts.accounts[22] = acc
  216. cnQuota.result = &CNProviderQuotaProbeResult{Success: true}
  217. fetcher.Fetch(context.Background(), 22)
  218. require.Equal(t, 1, accounts.calls)
  219. require.Same(t, acc, cnQuota.lastAccount)
  220. require.Equal(t, 1, cnQuota.calls)
  221. })
  222. t.Run("cn payg", func(t *testing.T) {
  223. fetcher, _, _, cnBalance, accounts := newQuotaFetcherTestSetup(t)
  224. acc := &Account{ID: 23, Platform: domain.PlatformDeepseek, Credentials: map[string]any{"account_mode": AccountModePayG}}
  225. accounts.accounts[23] = acc
  226. cnBalance.result = &CNProviderBalanceResult{Success: true, Available: true, Balance: 1, Currency: "CNY"}
  227. fetcher.Fetch(context.Background(), 23)
  228. require.Equal(t, 1, accounts.calls)
  229. require.Same(t, acc, cnBalance.lastAccount)
  230. require.Equal(t, 1, cnBalance.calls)
  231. })
  232. }
  233. // --- 失败路径(Fetch 永不返回 error) ---
  234. func TestQuotaFetcher_AccountMissingYieldsLinkedAccountSnapshot(t *testing.T) {
  235. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  236. accounts.err = errors.New("not found")
  237. snapshot := fetcher.Fetch(context.Background(), 404)
  238. require.False(t, snapshot.Success)
  239. require.Equal(t, "linked account not found", snapshot.Error)
  240. require.Equal(t, 0, usage.getCalls()) // 未走到数据源
  241. }
  242. func TestQuotaFetcher_UsageAuthErrorMarksCredentialInvalid(t *testing.T) {
  243. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  244. accounts.accounts[3] = &Account{ID: 3, Platform: domain.PlatformOpenAI}
  245. usage.err = errors.New("API returned 401: unauthorized")
  246. snapshot := fetcher.Fetch(context.Background(), 3)
  247. require.False(t, snapshot.Success)
  248. require.True(t, snapshot.CredentialInvalid)
  249. require.Contains(t, snapshot.Error, "401")
  250. }
  251. // 值通道失败:antigravity/grok 等平台 err==nil 但错误降级在 UsageInfo 字段里,
  252. // 必须识别为失败快照,否则会被误判为 operational。
  253. func TestQuotaFetcher_UsageValueChannelFailureYieldsFailureSnapshot(t *testing.T) {
  254. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  255. // 凭据失效(401 语义)→ failed。
  256. accounts.accounts[3] = &Account{ID: 3, Platform: domain.PlatformAnthropic}
  257. usage.usage = &UsageInfo{Error: "usage API error: HTTP 401", ErrorCode: errorCodeUnauthenticated, NeedsReauth: true}
  258. snapshot := fetcher.Fetch(context.Background(), 3)
  259. require.False(t, snapshot.Success)
  260. require.True(t, snapshot.CredentialInvalid)
  261. require.Contains(t, snapshot.Error, "401")
  262. require.Equal(t, MonitorStatusFailed, deriveQuotaCheckResult(snapshot, "quota", time.Now()).Status)
  263. // 限流等非凭据失败 → error(而非 operational)。
  264. accounts.accounts[13] = &Account{ID: 13, Platform: domain.PlatformAnthropic}
  265. usage.usage = &UsageInfo{Error: "usage API error: HTTP 429", ErrorCode: errorCodeRateLimited}
  266. snapshot = fetcher.Fetch(context.Background(), 13)
  267. require.False(t, snapshot.Success)
  268. require.False(t, snapshot.CredentialInvalid)
  269. require.Contains(t, snapshot.Error, "429")
  270. require.Equal(t, MonitorStatusError, deriveQuotaCheckResult(snapshot, "quota", time.Now()).Status)
  271. // grok 已知未知态(尚未观测到计费/限流头)不算失败。
  272. accounts.accounts[14] = &Account{ID: 14, Platform: domain.PlatformGrok}
  273. usage.usage = &UsageInfo{ErrorCode: "quota_unknown", Error: "Grok quota is unknown until billing is probed"}
  274. snapshot = fetcher.Fetch(context.Background(), 14)
  275. require.True(t, snapshot.Success)
  276. require.Empty(t, snapshot.Error)
  277. require.Empty(t, snapshot.Tiers)
  278. require.Equal(t, MonitorStatusOperational, deriveQuotaCheckResult(snapshot, "quota", time.Now()).Status)
  279. }
  280. func TestUsageFailureInfo_ClassificationMatrix(t *testing.T) {
  281. cases := []struct {
  282. name string
  283. usage *UsageInfo
  284. failed bool
  285. credentialInvalid bool
  286. msg string
  287. }{
  288. {name: "nil usage", usage: nil},
  289. {name: "healthy empty", usage: &UsageInfo{}},
  290. {name: "error text only", usage: &UsageInfo{Error: "boom"}, failed: true, msg: "boom"},
  291. {name: "needs reauth", usage: &UsageInfo{NeedsReauth: true}, failed: true, credentialInvalid: true, msg: "usage fetch failed"},
  292. {name: "banned", usage: &UsageInfo{IsBanned: true}, failed: true, credentialInvalid: true, msg: "usage fetch failed"},
  293. {name: "forbidden with reason", usage: &UsageInfo{IsForbidden: true, ForbiddenReason: "usage limited"}, failed: true, credentialInvalid: true, msg: "usage limited"},
  294. {name: "error code unauthenticated", usage: &UsageInfo{ErrorCode: errorCodeUnauthenticated}, failed: true, credentialInvalid: true, msg: errorCodeUnauthenticated},
  295. {name: "error code forbidden", usage: &UsageInfo{ErrorCode: errorCodeForbidden}, failed: true, credentialInvalid: true, msg: errorCodeForbidden},
  296. {name: "error code rate limited", usage: &UsageInfo{ErrorCode: errorCodeRateLimited}, failed: true, msg: errorCodeRateLimited},
  297. {name: "error code network error", usage: &UsageInfo{ErrorCode: errorCodeNetworkError}, failed: true, msg: errorCodeNetworkError},
  298. {name: "grok quota unknown exempted", usage: &UsageInfo{ErrorCode: "quota_unknown", Error: "Grok quota is unknown until billing is probed"}},
  299. }
  300. for _, tc := range cases {
  301. t.Run(tc.name, func(t *testing.T) {
  302. failed, credentialInvalid, msg := usageFailureInfo(tc.usage)
  303. require.Equal(t, tc.failed, failed)
  304. require.Equal(t, tc.credentialInvalid, credentialInvalid)
  305. if tc.msg != "" {
  306. require.Equal(t, tc.msg, msg)
  307. }
  308. })
  309. }
  310. }
  311. // 凭据失效只认 401/403(与 fetchCNBalance 口径一致):CN quota 服务的
  312. // CredentialValid 仅成功路径置 true,500/429/智谱业务错误须推导为 error 而非 failed。
  313. func TestQuotaFetcher_CNQuotaCredentialInvalidByStatusCode(t *testing.T) {
  314. cases := []struct {
  315. name string
  316. accountID int64
  317. statusCode int
  318. credentialBad bool
  319. expectedStatus string
  320. }{
  321. {name: "401 unauthorized", accountID: 5, statusCode: 401, credentialBad: true, expectedStatus: MonitorStatusFailed},
  322. {name: "403 forbidden", accountID: 15, statusCode: 403, credentialBad: true, expectedStatus: MonitorStatusFailed},
  323. {name: "500 server error", accountID: 16, statusCode: 500, expectedStatus: MonitorStatusError},
  324. {name: "429 rate limited", accountID: 17, statusCode: 429, expectedStatus: MonitorStatusError},
  325. // 智谱 2xx 但业务级失败:StatusCode=200,非凭据问题。
  326. {name: "200 business error", accountID: 18, statusCode: 200, expectedStatus: MonitorStatusError},
  327. }
  328. for _, tc := range cases {
  329. t.Run(tc.name, func(t *testing.T) {
  330. fetcher, _, cnQuota, _, accounts := newQuotaFetcherTestSetup(t)
  331. accounts.accounts[tc.accountID] = &Account{
  332. ID: tc.accountID,
  333. Platform: domain.PlatformZhipu,
  334. Credentials: map[string]any{"account_mode": AccountModeCoding},
  335. }
  336. cnQuota.result = &CNProviderQuotaProbeResult{
  337. Success: false,
  338. StatusCode: tc.statusCode,
  339. Error: "api key expired",
  340. }
  341. snapshot := fetcher.Fetch(context.Background(), tc.accountID)
  342. require.False(t, snapshot.Success)
  343. require.Equal(t, tc.credentialBad, snapshot.CredentialInvalid)
  344. require.Equal(t, tc.expectedStatus, deriveQuotaCheckResult(snapshot, "quota", time.Now()).Status)
  345. })
  346. }
  347. }
  348. func TestQuotaFetcher_CNBalanceHTTP403MarksCredentialInvalid(t *testing.T) {
  349. fetcher, _, _, cnBalance, accounts := newQuotaFetcherTestSetup(t)
  350. accounts.accounts[6] = &Account{ID: 6, Platform: domain.PlatformKimi}
  351. cnBalance.result = &CNProviderBalanceResult{Success: false, StatusCode: 403, Error: "forbidden"}
  352. snapshot := fetcher.Fetch(context.Background(), 6)
  353. require.False(t, snapshot.Success)
  354. require.True(t, snapshot.CredentialInvalid)
  355. }
  356. // 余额告警口径与账号停调(CNProviderBalanceCheckService.checkOne)一致:
  357. // 上游标记不可用或全部币种低于阈值 → BalanceLow → degraded;任一币种达标即健康。
  358. func TestQuotaFetcher_CNBalanceLowMarksDegraded(t *testing.T) {
  359. cases := []struct {
  360. name string
  361. accountID int64
  362. result *CNProviderBalanceResult
  363. balanceLow bool
  364. wantStatus string
  365. wantMessage string
  366. }{
  367. {
  368. // 审查例:余额 5/阈值 10 的账号调度器已停调,监控不能仍绿灯。
  369. name: "balance below threshold",
  370. accountID: 21,
  371. result: &CNProviderBalanceResult{Success: true, Available: true, Balance: 5, Currency: "CNY"},
  372. balanceLow: true,
  373. wantStatus: MonitorStatusDegraded, wantMessage: "balance low: 5 CNY",
  374. },
  375. {
  376. name: "upstream marked unavailable",
  377. accountID: 22,
  378. result: &CNProviderBalanceResult{Success: true, Available: false, Balance: 20, Currency: "CNY"},
  379. balanceLow: true,
  380. wantStatus: MonitorStatusDegraded, wantMessage: "balance low: 20 CNY",
  381. },
  382. {
  383. // deepseek 双币种:任一币种(USD 20)达标即健康。
  384. name: "any currency above threshold is healthy",
  385. accountID: 23,
  386. result: &CNProviderBalanceResult{
  387. Success: true, Available: true, Balance: 5, Currency: "CNY",
  388. Balances: []CNProviderBalanceEntry{{Currency: "CNY", Balance: 5}, {Currency: "USD", Balance: 20}},
  389. },
  390. wantStatus: MonitorStatusOperational,
  391. },
  392. {
  393. name: "single currency above threshold",
  394. accountID: 24,
  395. result: &CNProviderBalanceResult{Success: true, Available: true, Balance: 20, Currency: "CNY"},
  396. wantStatus: MonitorStatusOperational,
  397. },
  398. }
  399. for _, tc := range cases {
  400. t.Run(tc.name, func(t *testing.T) {
  401. fetcher, _, _, cnBalance, accounts := newQuotaFetcherTestSetup(t)
  402. fetcher.balanceThreshold = 10
  403. accounts.accounts[tc.accountID] = &Account{
  404. ID: tc.accountID,
  405. Platform: domain.PlatformKimi,
  406. Credentials: map[string]any{"account_mode": AccountModePayG},
  407. }
  408. cnBalance.result = tc.result
  409. snapshot := fetcher.Fetch(context.Background(), tc.accountID)
  410. require.True(t, snapshot.Success)
  411. require.Equal(t, tc.balanceLow, snapshot.BalanceLow)
  412. res := deriveQuotaCheckResult(snapshot, "quota", time.Now())
  413. require.Equal(t, tc.wantStatus, res.Status)
  414. if tc.wantMessage != "" {
  415. require.Contains(t, res.Message, tc.wantMessage)
  416. } else {
  417. require.Empty(t, res.Message)
  418. }
  419. })
  420. }
  421. }
  422. func TestNewChannelMonitorQuotaFetcher_ThresholdFromConfig(t *testing.T) {
  423. require.InDelta(t, 0.5, NewChannelMonitorQuotaFetcher(nil, nil, nil, nil, nil).balanceThreshold, 0.0001)
  424. cfg10 := &config.Config{Gateway: config.GatewayConfig{CNProviders: config.GatewayCNProvidersConfig{BalanceThreshold: 10}}}
  425. require.InDelta(t, 10, NewChannelMonitorQuotaFetcher(nil, nil, nil, nil, cfg10).balanceThreshold, 0.0001)
  426. // 非正值(含显式 0)回退默认,避免 0 阈值下「余额=0 也不告警」。
  427. cfg0 := &config.Config{Gateway: config.GatewayConfig{CNProviders: config.GatewayCNProvidersConfig{BalanceThreshold: 0}}}
  428. require.InDelta(t, 0.5, NewChannelMonitorQuotaFetcher(nil, nil, nil, nil, cfg0).balanceThreshold, 0.0001)
  429. }
  430. func TestQuotaFetcher_NilDependenciesProduceErrorSnapshots(t *testing.T) {
  431. // fetcher 本体为 nil:直接降级为错误快照,不 panic。
  432. var nilFetcher *ChannelMonitorQuotaFetcher
  433. snapshot := nilFetcher.Fetch(context.Background(), 1)
  434. require.False(t, snapshot.Success)
  435. require.Equal(t, "quota fetcher is not configured", snapshot.Error)
  436. // 数据源缺失:账号能加载,但对应服务未注入。
  437. fetcher, _, _, _, accounts := newQuotaFetcherTestSetup(t)
  438. fetcher.usage = nil
  439. accounts.accounts[2] = &Account{ID: 2, Platform: domain.PlatformOpenAI}
  440. snapshot = fetcher.Fetch(context.Background(), 2)
  441. require.False(t, snapshot.Success)
  442. require.Contains(t, snapshot.Error, "not configured")
  443. }
  444. // --- TTL 缓存 ---
  445. func TestQuotaFetcher_CachesSuccessSnapshotPerAccount(t *testing.T) {
  446. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  447. accounts.accounts[8] = &Account{ID: 8, Platform: domain.PlatformOpenAI}
  448. usage.usage = &UsageInfo{FiveHour: &UsageProgress{Utilization: 10}}
  449. for i := 0; i < 3; i++ {
  450. snapshot := fetcher.Fetch(context.Background(), 8)
  451. require.True(t, snapshot.Success)
  452. }
  453. require.Equal(t, 1, usage.getCalls(), "success snapshots should be served from cache")
  454. // 缓存过期后重新拉取。
  455. fetcher.mu.Lock()
  456. entry := fetcher.cache[8]
  457. entry.expiry = time.Now().Add(-time.Second)
  458. fetcher.cache[8] = entry
  459. fetcher.mu.Unlock()
  460. _ = fetcher.Fetch(context.Background(), 8)
  461. require.Equal(t, 2, usage.getCalls())
  462. }
  463. func TestQuotaFetcher_CachesFailureSnapshotWithShortTTL(t *testing.T) {
  464. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  465. accounts.accounts[4] = &Account{ID: 4, Platform: domain.PlatformOpenAI}
  466. usage.err = errors.New("boom")
  467. for i := 0; i < 2; i++ {
  468. snapshot := fetcher.Fetch(context.Background(), 4)
  469. require.False(t, snapshot.Success)
  470. }
  471. require.Equal(t, 1, usage.getCalls(), "failure snapshots should be served from the short negative cache")
  472. // 失败快照的 TTL 是负缓存时长(而非成功 TTL)。
  473. fetcher.mu.Lock()
  474. entry := fetcher.cache[4]
  475. require.WithinDuration(t, entry.snapshot.FetchedAt.Add(monitorQuotaErrorCacheTTL), entry.expiry, time.Second)
  476. entry.expiry = time.Now().Add(-time.Second)
  477. fetcher.cache[4] = entry
  478. fetcher.mu.Unlock()
  479. _ = fetcher.Fetch(context.Background(), 4)
  480. require.Equal(t, 2, usage.getCalls(), "expired negative cache should refetch")
  481. }
  482. func TestQuotaFetcher_ConcurrentFetchesShareSingleFlight(t *testing.T) {
  483. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  484. accounts.accounts[12] = &Account{ID: 12, Platform: domain.PlatformOpenAI}
  485. usage.usage = &UsageInfo{FiveHour: &UsageProgress{Utilization: 10}}
  486. usage.block = make(chan struct{})
  487. var wg sync.WaitGroup
  488. snapshots := make([]*domain.MonitorQuotaSnapshot, 5)
  489. for i := 0; i < 5; i++ {
  490. wg.Add(1)
  491. go func(idx int) {
  492. defer wg.Done()
  493. snapshots[idx] = fetcher.Fetch(context.Background(), 12)
  494. }(i)
  495. }
  496. // 上游被 block 卡住时,5 个并发 Fetch 应只产生 1 次真实查询。
  497. require.Eventually(t, func() bool { return usage.getCalls() == 1 },
  498. 5*time.Second, 10*time.Millisecond)
  499. close(usage.block)
  500. wg.Wait()
  501. for _, snapshot := range snapshots {
  502. require.NotNil(t, snapshot)
  503. require.True(t, snapshot.Success)
  504. }
  505. require.Equal(t, 1, usage.getCalls())
  506. // 成功快照已缓存:再取一次仍不打上游。
  507. _ = fetcher.Fetch(context.Background(), 12)
  508. require.Equal(t, 1, usage.getCalls())
  509. }
  510. // 上面那条并发用例只能以极低概率撞上真正的缺陷窗口(实测约 300 次一次),
  511. // 所以这里直接钉住窗口本身。
  512. //
  513. // 窗口在 Fetch 顶部的 cachedSnapshot 与 flight.DoChan 之间:一个 goroutine 判定
  514. // 缓存未命中之后、真正入队之前,另一个 goroutine 的 flight 可能已经跑完、写好缓存
  515. // 并被摘掉 key,于是前者不会并入那次飞行,而是另起一次新的,对同一账号重复打上游。
  516. // fetchShared 就是 flight 的执行体,直接调用它等价于「已经越过顶层缓存判定、拿到了
  517. // 属于自己的那次飞行」这个状态。
  518. func TestQuotaFetcher_SharedFetchCacheRecheckAvoidsDuplicateUpstream(t *testing.T) {
  519. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  520. accounts.accounts[13] = &Account{ID: 13, Platform: domain.PlatformOpenAI}
  521. usage.usage = &UsageInfo{FiveHour: &UsageProgress{Utilization: 10}}
  522. first := fetcher.Fetch(context.Background(), 13)
  523. require.True(t, first.Success)
  524. require.Equal(t, 1, usage.getCalls())
  525. shared := fetcher.fetchShared(13)
  526. require.Equal(t, 1, usage.getCalls(), "重查缓存后不得再打一次上游")
  527. require.Same(t, first, shared, "应原样返回已缓存的快照")
  528. }
  529. // 重查不能变成无条件短路:缓存过期后同一个执行体必须照常回源,
  530. // 否则快照会永远停在第一次的值上。
  531. func TestQuotaFetcher_SharedFetchStillRefetchesAfterCacheExpiry(t *testing.T) {
  532. fetcher, usage, _, _, accounts := newQuotaFetcherTestSetup(t)
  533. accounts.accounts[14] = &Account{ID: 14, Platform: domain.PlatformOpenAI}
  534. usage.usage = &UsageInfo{FiveHour: &UsageProgress{Utilization: 10}}
  535. require.True(t, fetcher.Fetch(context.Background(), 14).Success)
  536. require.Equal(t, 1, usage.getCalls())
  537. fetcher.mu.Lock()
  538. entry := fetcher.cache[14]
  539. entry.expiry = time.Now().Add(-time.Second)
  540. fetcher.cache[14] = entry
  541. fetcher.mu.Unlock()
  542. require.True(t, fetcher.fetchShared(14).Success)
  543. require.Equal(t, 2, usage.getCalls(), "缓存过期后必须回源")
  544. }
  545. // --- UsageInfo → tiers 归一 ---
  546. func TestUsageQuotaTiers_MapsAllWindowKinds(t *testing.T) {
  547. limit := int64(1000)
  548. remaining := int64(400)
  549. resetUnix := int64(1777283883)
  550. usage := &UsageInfo{
  551. FiveHour: &UsageProgress{Utilization: 50},
  552. SevenDay: &UsageProgress{Utilization: 60},
  553. SevenDaySonnet: &UsageProgress{Utilization: 70},
  554. SevenDayFable: &UsageProgress{Utilization: 80},
  555. ThirtyDay: &UsageProgress{Utilization: 20},
  556. GeminiSharedDaily: &UsageProgress{Utilization: 11},
  557. GeminiProDaily: &UsageProgress{Utilization: 22},
  558. GeminiFlashDaily: &UsageProgress{Utilization: 33},
  559. GrokRequestQuota: &xai.QuotaWindow{Limit: &limit, Remaining: &remaining, ResetUnix: &resetUnix},
  560. GrokTokenQuota: &xai.QuotaWindow{Limit: &limit, Remaining: &remaining, ResetAt: "2026-08-19T00:00:00Z"},
  561. AntigravityQuota: map[string]*AntigravityModelQuota{
  562. "gemini-3-pro": {Utilization: 45},
  563. "gemini-3-flash": {Utilization: 55},
  564. },
  565. }
  566. tiers := usageQuotaTiers(usage)
  567. // 5h/7d/7d-sonnet/7d-fable/30d + gemini×3 + grok×2 + antigravity×2
  568. require.Len(t, tiers, 12)
  569. byKey := make(map[string]domain.MonitorQuotaTier, len(tiers))
  570. for _, tier := range tiers {
  571. key := tier.Window
  572. if tier.Label != "" {
  573. key = tier.Window + "/" + tier.Label
  574. }
  575. byKey[key] = tier
  576. }
  577. require.Contains(t, byKey, "5h")
  578. require.Contains(t, byKey, "7d")
  579. require.Contains(t, byKey, "7d-sonnet")
  580. require.Contains(t, byKey, "7d-fable")
  581. require.Contains(t, byKey, "30d")
  582. require.Contains(t, byKey, "daily/shared")
  583. require.Contains(t, byKey, "daily/pro")
  584. require.Contains(t, byKey, "daily/flash")
  585. require.Contains(t, byKey, "daily/requests")
  586. require.Contains(t, byKey, "daily/tokens")
  587. require.Contains(t, byKey, "total/gemini-3-pro")
  588. require.Contains(t, byKey, "total/gemini-3-flash")
  589. // grok requests 窗口:used = limit - remaining,百分比 60%。
  590. requests := byKey["daily/requests"]
  591. require.Equal(t, float64(600), requests.Used)
  592. require.Equal(t, float64(1000), requests.Limit)
  593. require.InDelta(t, 60.0, requests.UsedPercent, 0.001)
  594. require.NotEmpty(t, requests.ResetAt, "ResetUnix should fall back to RFC3339")
  595. tokens := byKey["daily/tokens"]
  596. require.Equal(t, "2026-08-19T00:00:00Z", tokens.ResetAt)
  597. }
  598. func TestUsageQuotaTiers_NilAndEmptyInputs(t *testing.T) {
  599. require.Nil(t, usageQuotaTiers(nil))
  600. require.Nil(t, usageQuotaTiers(&UsageInfo{}))
  601. // Grok 窗口 limit<=0 时跳过,避免除零。
  602. var zero int64
  603. tiers := usageQuotaTiers(&UsageInfo{
  604. GrokRequestQuota: &xai.QuotaWindow{Limit: &zero, Remaining: &zero},
  605. })
  606. require.Nil(t, tiers)
  607. }
  608. // --- 状态推导 ---
  609. func TestDeriveQuotaCheckResult_StatusMatrix(t *testing.T) {
  610. now := time.Now()
  611. healthy := &domain.MonitorQuotaSnapshot{Success: true, Tiers: []domain.MonitorQuotaTier{{Window: "5h", UsedPercent: 40}}}
  612. res := deriveQuotaCheckResult(healthy, "quota", now)
  613. require.Equal(t, MonitorStatusOperational, res.Status)
  614. require.Equal(t, "quota", res.Model)
  615. require.Empty(t, res.Message)
  616. highUsage := &domain.MonitorQuotaSnapshot{Success: true, Tiers: []domain.MonitorQuotaTier{
  617. {Window: "5h", UsedPercent: 30},
  618. {Window: "daily", Label: "pro", UsedPercent: 95},
  619. }}
  620. res = deriveQuotaCheckResult(highUsage, "quota", now)
  621. require.Equal(t, MonitorStatusDegraded, res.Status)
  622. require.Contains(t, res.Message, "pro/daily")
  623. require.Contains(t, res.Message, "95.0%")
  624. balance := -0.5
  625. lowBalance := &domain.MonitorQuotaSnapshot{Success: true, BalanceLow: true, Balance: &balance, Currency: "CNY"}
  626. res = deriveQuotaCheckResult(lowBalance, "quota", now)
  627. require.Equal(t, MonitorStatusDegraded, res.Status)
  628. require.Contains(t, res.Message, "balance low")
  629. invalid := &domain.MonitorQuotaSnapshot{Success: false, CredentialInvalid: true, Error: "401 unauthorized"}
  630. res = deriveQuotaCheckResult(invalid, "quota", now)
  631. require.Equal(t, MonitorStatusFailed, res.Status)
  632. unlinked := &domain.MonitorQuotaSnapshot{Success: false, Error: "linked account not found"}
  633. res = deriveQuotaCheckResult(unlinked, "quota", now)
  634. require.Equal(t, MonitorStatusDegraded, res.Status)
  635. other := &domain.MonitorQuotaSnapshot{Success: false, Error: "connection refused"}
  636. res = deriveQuotaCheckResult(other, "quota", now)
  637. require.Equal(t, MonitorStatusError, res.Status)
  638. require.Equal(t, "connection refused", res.Message)
  639. res = deriveQuotaCheckResult(nil, "quota", now)
  640. require.Equal(t, MonitorStatusError, res.Status)
  641. }