openai_oauth_handler_reset_quota_test.go 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492
  1. //go:build unit
  2. package admin
  3. import (
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "net/http"
  8. "net/http/httptest"
  9. "testing"
  10. "github.com/gin-gonic/gin"
  11. "github.com/stretchr/testify/require"
  12. "github.com/Wei-Shaw/sub2api/internal/service"
  13. )
  14. type openAIQuotaWorkflowStub struct {
  15. resetResult *service.OpenAIQuotaResetResult
  16. resetErr error
  17. queryResult *service.OpenAIQuotaUsage
  18. queryErr error
  19. cacheErr error
  20. resetCalls int
  21. queryCalls int
  22. cacheCalls int
  23. queryCtxErr error
  24. cacheCtxErr error
  25. }
  26. func (s *openAIQuotaWorkflowStub) ResetCredit(context.Context, int64) (*service.OpenAIQuotaResetResult, error) {
  27. s.resetCalls++
  28. return s.resetResult, s.resetErr
  29. }
  30. func (s *openAIQuotaWorkflowStub) QueryUsage(ctx context.Context, _ int64) (*service.OpenAIQuotaUsage, error) {
  31. s.queryCalls++
  32. s.queryCtxErr = ctx.Err()
  33. return s.queryResult, s.queryErr
  34. }
  35. func (s *openAIQuotaWorkflowStub) CacheResetCreditsSnapshot(ctx context.Context, _ int64, _ *service.OpenAIRateLimitResetCredits) error {
  36. s.cacheCalls++
  37. s.cacheCtxErr = ctx.Err()
  38. return s.cacheErr
  39. }
  40. func (s *openAIQuotaWorkflowStub) CachePostResetSnapshot(ctx context.Context, _ int64, _ *service.OpenAIQuotaUsage) error {
  41. s.cacheCalls++
  42. s.cacheCtxErr = ctx.Err()
  43. return s.cacheErr
  44. }
  45. type openAIAccountStateRecovererStub struct {
  46. err error
  47. calls int
  48. accountID int64
  49. lastOptions service.AccountRecoveryOptions
  50. lastCtxErr error
  51. }
  52. func (s *openAIAccountStateRecovererStub) RecoverAccountState(ctx context.Context, accountID int64, options service.AccountRecoveryOptions) (*service.SuccessfulTestRecoveryResult, error) {
  53. s.calls++
  54. s.accountID = accountID
  55. s.lastOptions = options
  56. s.lastCtxErr = ctx.Err()
  57. return &service.SuccessfulTestRecoveryResult{}, s.err
  58. }
  59. type openAIResetAdminServiceStub struct {
  60. service.AdminService
  61. account *service.Account
  62. err error
  63. calls int
  64. }
  65. func (s *openAIResetAdminServiceStub) GetAccount(context.Context, int64) (*service.Account, error) {
  66. s.calls++
  67. return s.account, s.err
  68. }
  69. type openAIQuotaResetEnvelope struct {
  70. Code int `json:"code"`
  71. Data openAIQuotaResetResponse `json:"data"`
  72. }
  73. type openAIQuotaRefreshEnvelope struct {
  74. Code int `json:"code"`
  75. Data openAIQuotaRefreshResponse `json:"data"`
  76. }
  77. func performOpenAIQuotaResetRequest(t *testing.T, handler *OpenAIOAuthHandler) (int, openAIQuotaResetEnvelope) {
  78. t.Helper()
  79. return performOpenAIQuotaResetRequestWithContext(t, handler, nil)
  80. }
  81. // performOpenAIQuotaResetRequestWithContext drives the reset endpoint, optionally
  82. // with an already-canceled request context (client disconnect simulation).
  83. func performOpenAIQuotaResetRequestWithContext(t *testing.T, handler *OpenAIOAuthHandler, ctx context.Context) (int, openAIQuotaResetEnvelope) {
  84. t.Helper()
  85. gin.SetMode(gin.TestMode)
  86. router := gin.New()
  87. router.POST("/api/v1/admin/openai/accounts/:id/reset-quota", handler.ResetQuota)
  88. recorder := httptest.NewRecorder()
  89. request := httptest.NewRequest(http.MethodPost, "/api/v1/admin/openai/accounts/42/reset-quota", nil)
  90. if ctx != nil {
  91. request = request.WithContext(ctx)
  92. }
  93. router.ServeHTTP(recorder, request)
  94. var envelope openAIQuotaResetEnvelope
  95. require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &envelope))
  96. return recorder.Code, envelope
  97. }
  98. func performOpenAIQuotaRefreshRequest(t *testing.T, handler *OpenAIOAuthHandler) (int, openAIQuotaRefreshEnvelope) {
  99. t.Helper()
  100. gin.SetMode(gin.TestMode)
  101. router := gin.New()
  102. router.POST("/api/v1/admin/openai/accounts/:id/quota/refresh", handler.RefreshQuota)
  103. recorder := httptest.NewRecorder()
  104. request := httptest.NewRequest(http.MethodPost, "/api/v1/admin/openai/accounts/42/quota/refresh", nil)
  105. router.ServeHTTP(recorder, request)
  106. var envelope openAIQuotaRefreshEnvelope
  107. require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &envelope))
  108. return recorder.Code, envelope
  109. }
  110. func successfulOpenAIQuotaWorkflowStub() *openAIQuotaWorkflowStub {
  111. return &openAIQuotaWorkflowStub{
  112. resetResult: &service.OpenAIQuotaResetResult{
  113. Code: "success",
  114. WindowsReset: 1,
  115. },
  116. queryResult: &service.OpenAIQuotaUsage{
  117. FetchedAt: 123,
  118. RateLimitResetCredits: &service.OpenAIRateLimitResetCredits{
  119. AvailableCount: 0,
  120. Credits: []service.OpenAIRateLimitResetCreditDetail{},
  121. },
  122. },
  123. }
  124. }
  125. func recoveredAccountStub() *openAIResetAdminServiceStub {
  126. return &openAIResetAdminServiceStub{account: &service.Account{
  127. ID: 42,
  128. Name: "recovered",
  129. Platform: service.PlatformOpenAI,
  130. Type: service.AccountTypeOAuth,
  131. Status: service.StatusActive,
  132. Schedulable: false,
  133. }}
  134. }
  135. func TestOpenAIResetQuota_ResetFailureStopsWorkflow(t *testing.T) {
  136. quota := &openAIQuotaWorkflowStub{resetErr: errors.New("upstream reset failed")}
  137. recoverer := &openAIAccountStateRecovererStub{}
  138. handler := &OpenAIOAuthHandler{
  139. adminService: &openAIResetAdminServiceStub{},
  140. quotaService: quota,
  141. rateLimitService: recoverer,
  142. }
  143. status, _ := performOpenAIQuotaResetRequest(t, handler)
  144. require.Equal(t, http.StatusInternalServerError, status)
  145. require.Equal(t, 1, quota.resetCalls)
  146. require.Zero(t, quota.queryCalls)
  147. require.Zero(t, quota.cacheCalls)
  148. require.Zero(t, recoverer.calls)
  149. }
  150. // Account-state recovery is the reason the credit was spent (#3672 / #3740), so it
  151. // must run before — and independently of — the reset-credit display cache.
  152. func TestOpenAIResetQuota_RecoversAccountStateBeforeRefreshingCache(t *testing.T) {
  153. quota := successfulOpenAIQuotaWorkflowStub()
  154. recoverer := &openAIAccountStateRecovererStub{}
  155. adminService := recoveredAccountStub()
  156. handler := &OpenAIOAuthHandler{
  157. adminService: adminService,
  158. quotaService: quota,
  159. rateLimitService: recoverer,
  160. }
  161. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  162. require.Equal(t, http.StatusOK, status)
  163. require.Empty(t, envelope.Data.WarningCode)
  164. require.True(t, envelope.Data.AccountStateRecovered)
  165. require.True(t, envelope.Data.CacheRefreshed)
  166. require.NotNil(t, envelope.Data.Quota)
  167. require.NotNil(t, envelope.Data.Account)
  168. require.Equal(t, int64(42), envelope.Data.Account.ID)
  169. require.False(t, envelope.Data.Account.Schedulable, "manual scheduling switch must not be flipped")
  170. require.Equal(t, int64(42), recoverer.accountID)
  171. require.True(t, recoverer.lastOptions.InvalidateToken)
  172. require.Equal(t, 1, quota.resetCalls)
  173. require.Equal(t, 1, quota.queryCalls)
  174. require.Equal(t, 1, quota.cacheCalls)
  175. require.Equal(t, 1, recoverer.calls)
  176. require.Equal(t, 1, adminService.calls)
  177. }
  178. func TestOpenAIResetQuota_RecoveryFailureStopsWorkflow(t *testing.T) {
  179. quota := successfulOpenAIQuotaWorkflowStub()
  180. recoverer := &openAIAccountStateRecovererStub{err: errors.New("recovery failed")}
  181. adminService := recoveredAccountStub()
  182. handler := &OpenAIOAuthHandler{
  183. adminService: adminService,
  184. quotaService: quota,
  185. rateLimitService: recoverer,
  186. }
  187. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  188. require.Equal(t, http.StatusOK, status)
  189. require.Equal(t, service.OpenAIQuotaResetWarningAccountRecoveryFailed, envelope.Data.WarningCode)
  190. require.False(t, envelope.Data.AccountStateRecovered)
  191. require.False(t, envelope.Data.CacheRefreshed)
  192. require.Nil(t, envelope.Data.Quota)
  193. require.Nil(t, envelope.Data.Account)
  194. require.Equal(t, 1, recoverer.calls)
  195. require.Zero(t, quota.queryCalls)
  196. require.Zero(t, quota.cacheCalls)
  197. require.Zero(t, adminService.calls)
  198. }
  199. func TestOpenAIResetQuota_MissingRecovererReportsRecoveryFailure(t *testing.T) {
  200. quota := successfulOpenAIQuotaWorkflowStub()
  201. adminService := recoveredAccountStub()
  202. handler := &OpenAIOAuthHandler{
  203. adminService: adminService,
  204. quotaService: quota,
  205. }
  206. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  207. require.Equal(t, http.StatusOK, status)
  208. require.Equal(t, service.OpenAIQuotaResetWarningAccountRecoveryFailed, envelope.Data.WarningCode)
  209. require.False(t, envelope.Data.AccountStateRecovered)
  210. require.Zero(t, quota.queryCalls)
  211. require.Zero(t, adminService.calls)
  212. }
  213. // A failed cache refresh must not hide the recovered account row: the operator
  214. // still needs the list to drop the stale rate-limit badge.
  215. func TestOpenAIResetQuota_QueryFailureStillRecoversAndReturnsAccount(t *testing.T) {
  216. quota := successfulOpenAIQuotaWorkflowStub()
  217. quota.queryResult = nil
  218. quota.queryErr = errors.New("upstream query failed")
  219. recoverer := &openAIAccountStateRecovererStub{}
  220. adminService := recoveredAccountStub()
  221. handler := &OpenAIOAuthHandler{
  222. adminService: adminService,
  223. quotaService: quota,
  224. rateLimitService: recoverer,
  225. }
  226. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  227. require.Equal(t, http.StatusOK, status)
  228. require.Equal(t, service.OpenAIQuotaResetWarningCacheRefreshFailed, envelope.Data.WarningCode)
  229. require.True(t, envelope.Data.AccountStateRecovered)
  230. require.False(t, envelope.Data.CacheRefreshed)
  231. require.Nil(t, envelope.Data.Quota)
  232. require.NotNil(t, envelope.Data.Account)
  233. require.Equal(t, 1, quota.queryCalls)
  234. require.Zero(t, quota.cacheCalls)
  235. require.Equal(t, 1, recoverer.calls)
  236. require.Equal(t, 1, adminService.calls)
  237. }
  238. func TestOpenAIResetQuota_CacheFailureStillRecoversAndReturnsAccount(t *testing.T) {
  239. quota := successfulOpenAIQuotaWorkflowStub()
  240. quota.cacheErr = errors.New("cache write failed")
  241. recoverer := &openAIAccountStateRecovererStub{}
  242. adminService := recoveredAccountStub()
  243. handler := &OpenAIOAuthHandler{
  244. adminService: adminService,
  245. quotaService: quota,
  246. rateLimitService: recoverer,
  247. }
  248. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  249. require.Equal(t, http.StatusOK, status)
  250. require.Equal(t, service.OpenAIQuotaResetWarningCacheRefreshFailed, envelope.Data.WarningCode)
  251. require.True(t, envelope.Data.AccountStateRecovered)
  252. require.False(t, envelope.Data.CacheRefreshed)
  253. require.Nil(t, envelope.Data.Quota)
  254. require.NotNil(t, envelope.Data.Account)
  255. require.Equal(t, 1, quota.cacheCalls)
  256. require.Equal(t, 1, adminService.calls)
  257. }
  258. func TestOpenAIResetQuota_AccountRefreshFailureReportsRecoveredState(t *testing.T) {
  259. quota := successfulOpenAIQuotaWorkflowStub()
  260. recoverer := &openAIAccountStateRecovererStub{}
  261. adminService := &openAIResetAdminServiceStub{err: errors.New("account refresh failed")}
  262. handler := &OpenAIOAuthHandler{
  263. adminService: adminService,
  264. quotaService: quota,
  265. rateLimitService: recoverer,
  266. }
  267. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  268. require.Equal(t, http.StatusOK, status)
  269. require.Equal(t, service.OpenAIQuotaResetWarningAccountRefreshFailed, envelope.Data.WarningCode)
  270. require.True(t, envelope.Data.CacheRefreshed)
  271. require.True(t, envelope.Data.AccountStateRecovered)
  272. require.NotNil(t, envelope.Data.Quota)
  273. require.Nil(t, envelope.Data.Account)
  274. require.Equal(t, 1, adminService.calls)
  275. }
  276. // The first (most actionable) failure wins so the UI never downgrades a cache
  277. // problem into a cosmetic "could not reload the row" message.
  278. func TestOpenAIResetQuota_CacheAndAccountFailureKeepsFirstWarning(t *testing.T) {
  279. quota := successfulOpenAIQuotaWorkflowStub()
  280. quota.cacheErr = errors.New("cache write failed")
  281. recoverer := &openAIAccountStateRecovererStub{}
  282. adminService := &openAIResetAdminServiceStub{err: errors.New("account refresh failed")}
  283. handler := &OpenAIOAuthHandler{
  284. adminService: adminService,
  285. quotaService: quota,
  286. rateLimitService: recoverer,
  287. }
  288. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  289. require.Equal(t, http.StatusOK, status)
  290. require.Equal(t, service.OpenAIQuotaResetWarningCacheRefreshFailed, envelope.Data.WarningCode)
  291. require.True(t, envelope.Data.AccountStateRecovered)
  292. require.Nil(t, envelope.Data.Account)
  293. }
  294. // The credit is non-refundable once consumed, so post-reset bookkeeping must
  295. // survive a client disconnect instead of leaving the account rate-limited.
  296. func TestOpenAIResetQuota_PostProcessingSurvivesClientCancellation(t *testing.T) {
  297. quota := successfulOpenAIQuotaWorkflowStub()
  298. recoverer := &openAIAccountStateRecovererStub{}
  299. adminService := recoveredAccountStub()
  300. handler := &OpenAIOAuthHandler{
  301. adminService: adminService,
  302. quotaService: quota,
  303. rateLimitService: recoverer,
  304. }
  305. ctx, cancel := context.WithCancel(context.Background())
  306. cancel()
  307. status, _ := performOpenAIQuotaResetRequestWithContext(t, handler, ctx)
  308. require.Equal(t, http.StatusOK, status)
  309. require.Equal(t, 1, recoverer.calls)
  310. require.NoError(t, recoverer.lastCtxErr, "recovery must not inherit the canceled client context")
  311. require.NoError(t, quota.queryCtxErr)
  312. require.NoError(t, quota.cacheCtxErr)
  313. require.Equal(t, 1, adminService.calls)
  314. }
  315. func TestOpenAIRefreshQuota_PersistsSnapshot(t *testing.T) {
  316. quota := successfulOpenAIQuotaWorkflowStub()
  317. handler := &OpenAIOAuthHandler{
  318. adminService: &openAIResetAdminServiceStub{},
  319. quotaService: quota,
  320. }
  321. status, envelope := performOpenAIQuotaRefreshRequest(t, handler)
  322. require.Equal(t, http.StatusOK, status)
  323. require.True(t, envelope.Data.CachePersisted)
  324. require.Equal(t, int64(123), envelope.Data.FetchedAt)
  325. require.Equal(t, 1, quota.queryCalls)
  326. require.Equal(t, 1, quota.cacheCalls)
  327. require.Zero(t, quota.resetCalls)
  328. }
  329. // A rejected snapshot write must never discard the usage payload: otherwise the
  330. // card loses its credit count and the reset button stays disabled forever.
  331. func TestOpenAIRefreshQuota_PersistFailureStillReturnsUsage(t *testing.T) {
  332. quota := successfulOpenAIQuotaWorkflowStub()
  333. quota.queryResult = &service.OpenAIQuotaUsage{
  334. FetchedAt: 456,
  335. RateLimitResetCredits: &service.OpenAIRateLimitResetCredits{
  336. AvailableCount: 2,
  337. },
  338. }
  339. quota.cacheErr = errors.New("expiration details unavailable")
  340. handler := &OpenAIOAuthHandler{
  341. adminService: &openAIResetAdminServiceStub{},
  342. quotaService: quota,
  343. }
  344. status, envelope := performOpenAIQuotaRefreshRequest(t, handler)
  345. require.Equal(t, http.StatusOK, status)
  346. require.False(t, envelope.Data.CachePersisted)
  347. require.Equal(t, int64(456), envelope.Data.FetchedAt)
  348. require.NotNil(t, envelope.Data.RateLimitResetCredits)
  349. require.Equal(t, 2, envelope.Data.RateLimitResetCredits.AvailableCount)
  350. require.Equal(t, 1, quota.cacheCalls)
  351. }
  352. // An empty-but-successful upstream read must not be dereferenced blindly.
  353. func TestOpenAIQuotaEmptyUsageIsHandledWithoutPanic(t *testing.T) {
  354. t.Run("refresh reports an internal error", func(t *testing.T) {
  355. quota := successfulOpenAIQuotaWorkflowStub()
  356. quota.queryResult = nil
  357. handler := &OpenAIOAuthHandler{
  358. adminService: &openAIResetAdminServiceStub{},
  359. quotaService: quota,
  360. }
  361. status, envelope := performOpenAIQuotaRefreshRequest(t, handler)
  362. require.Equal(t, http.StatusInternalServerError, status)
  363. require.False(t, envelope.Data.CachePersisted)
  364. require.Zero(t, quota.cacheCalls)
  365. })
  366. t.Run("reset degrades to a cache warning", func(t *testing.T) {
  367. quota := successfulOpenAIQuotaWorkflowStub()
  368. quota.queryResult = nil
  369. recoverer := &openAIAccountStateRecovererStub{}
  370. adminService := recoveredAccountStub()
  371. handler := &OpenAIOAuthHandler{
  372. adminService: adminService,
  373. quotaService: quota,
  374. rateLimitService: recoverer,
  375. }
  376. status, envelope := performOpenAIQuotaResetRequest(t, handler)
  377. require.Equal(t, http.StatusOK, status)
  378. require.Equal(t, service.OpenAIQuotaResetWarningCacheRefreshFailed, envelope.Data.WarningCode)
  379. require.True(t, envelope.Data.AccountStateRecovered)
  380. require.NotNil(t, envelope.Data.Account)
  381. require.Zero(t, quota.cacheCalls)
  382. })
  383. }
  384. func TestOpenAIRefreshQuota_QueryFailureIsReported(t *testing.T) {
  385. quota := &openAIQuotaWorkflowStub{queryErr: errors.New("upstream query failed")}
  386. handler := &OpenAIOAuthHandler{
  387. adminService: &openAIResetAdminServiceStub{},
  388. quotaService: quota,
  389. }
  390. status, _ := performOpenAIQuotaRefreshRequest(t, handler)
  391. require.Equal(t, http.StatusInternalServerError, status)
  392. require.Equal(t, 1, quota.queryCalls)
  393. require.Zero(t, quota.cacheCalls)
  394. }
  395. // Storing a nil *Service in an interface field would make the capability guards
  396. // non-nil and panic on the first call; the constructor must keep them nil.
  397. func TestNewOpenAIOAuthHandlerKeepsNilQuotaCapabilitiesGuarded(t *testing.T) {
  398. gin.SetMode(gin.TestMode)
  399. handler := NewOpenAIOAuthHandler(nil, newStubAdminService(), nil, nil)
  400. require.Nil(t, handler.quotaService)
  401. require.Nil(t, handler.rateLimitService)
  402. router := gin.New()
  403. router.Use(gin.Recovery())
  404. router.GET("/openai/accounts/:id/quota", handler.QueryQuota)
  405. router.POST("/openai/accounts/:id/quota/refresh", handler.RefreshQuota)
  406. router.POST("/openai/accounts/:id/reset-quota", handler.ResetQuota)
  407. for _, tc := range []struct {
  408. method string
  409. path string
  410. }{
  411. {http.MethodGet, "/openai/accounts/42/quota"},
  412. {http.MethodPost, "/openai/accounts/42/quota/refresh"},
  413. {http.MethodPost, "/openai/accounts/42/reset-quota"},
  414. } {
  415. recorder := httptest.NewRecorder()
  416. router.ServeHTTP(recorder, httptest.NewRequest(tc.method, tc.path, nil))
  417. require.Equal(t, http.StatusBadRequest, recorder.Code, "%s %s", tc.method, tc.path)
  418. }
  419. }