account_repo_integration_test.go 65 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "database/sql"
  6. "strings"
  7. "testing"
  8. "time"
  9. dbent "github.com/Wei-Shaw/sub2api/ent"
  10. "github.com/Wei-Shaw/sub2api/ent/accountgroup"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  12. "github.com/Wei-Shaw/sub2api/internal/service"
  13. "github.com/stretchr/testify/require"
  14. "github.com/stretchr/testify/suite"
  15. )
  16. type AccountRepoSuite struct {
  17. suite.Suite
  18. ctx context.Context
  19. client *dbent.Client
  20. repo *accountRepository
  21. }
  22. type schedulerCacheRecorder struct {
  23. setAccounts []*service.Account
  24. deleteIDs []int64
  25. accounts map[int64]*service.Account
  26. setCtxErr error
  27. }
  28. func (s *schedulerCacheRecorder) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
  29. return nil, false, nil
  30. }
  31. func (s *schedulerCacheRecorder) CaptureBucketWriteToken(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
  32. return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
  33. }
  34. func (s *schedulerCacheRecorder) SetSnapshot(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) error {
  35. return nil
  36. }
  37. func (s *schedulerCacheRecorder) RetireBucket(ctx context.Context, bucket service.SchedulerBucket) error {
  38. return nil
  39. }
  40. func (s *schedulerCacheRecorder) ReopenBucket(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
  41. return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
  42. }
  43. func (s *schedulerCacheRecorder) TryAcquireGroupLifecycleLease(_ context.Context, groupID int64, _ time.Duration) (service.SchedulerGroupLifecycleLease, bool, error) {
  44. return service.SchedulerGroupLifecycleLease{GroupID: groupID, OwnerToken: "scheduler-cache-recorder"}, true, nil
  45. }
  46. func (s *schedulerCacheRecorder) ReleaseGroupLifecycleLease(context.Context, service.SchedulerGroupLifecycleLease) error {
  47. return nil
  48. }
  49. func (s *schedulerCacheRecorder) GetAccount(ctx context.Context, accountID int64) (*service.Account, error) {
  50. if s.accounts == nil {
  51. return nil, nil
  52. }
  53. return s.accounts[accountID], nil
  54. }
  55. func (s *schedulerCacheRecorder) SetAccount(ctx context.Context, account *service.Account) error {
  56. s.setCtxErr = ctx.Err()
  57. s.setAccounts = append(s.setAccounts, account)
  58. if s.accounts == nil {
  59. s.accounts = make(map[int64]*service.Account)
  60. }
  61. if account != nil {
  62. s.accounts[account.ID] = account
  63. }
  64. return nil
  65. }
  66. type failAtomicSchedulerOutboxSQLExecutor struct {
  67. sqlExecutor
  68. }
  69. func (e *failAtomicSchedulerOutboxSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
  70. if strings.Contains(query, "WITH updated AS") && strings.Contains(query, "INSERT INTO scheduler_outbox") && len(args) > 0 {
  71. args = append([]any(nil), args...)
  72. args[len(args)-1] = nil // event_type is NOT NULL; the whole statement must roll back.
  73. }
  74. return e.sqlExecutor.ExecContext(ctx, query, args...)
  75. }
  76. type cancelAfterAtomicMutationSQLExecutor struct {
  77. sqlExecutor
  78. cancel context.CancelFunc
  79. }
  80. func (e *cancelAfterAtomicMutationSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
  81. result, err := e.sqlExecutor.ExecContext(ctx, query, args...)
  82. if err == nil && strings.Contains(query, "WITH updated AS") && strings.Contains(query, "INSERT INTO scheduler_outbox") {
  83. e.cancel()
  84. }
  85. return result, err
  86. }
  87. func (s *schedulerCacheRecorder) DeleteAccount(ctx context.Context, accountID int64) error {
  88. s.deleteIDs = append(s.deleteIDs, accountID)
  89. if s.accounts != nil {
  90. delete(s.accounts, accountID)
  91. }
  92. return nil
  93. }
  94. func (s *schedulerCacheRecorder) UpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
  95. return nil
  96. }
  97. func (s *schedulerCacheRecorder) TryLockBucket(ctx context.Context, bucket service.SchedulerBucket, ttl time.Duration) (bool, error) {
  98. return true, nil
  99. }
  100. func (s *schedulerCacheRecorder) UnlockBucket(ctx context.Context, bucket service.SchedulerBucket) error {
  101. return nil
  102. }
  103. func (s *schedulerCacheRecorder) ListBuckets(ctx context.Context) ([]service.SchedulerBucket, error) {
  104. return nil, nil
  105. }
  106. func (s *schedulerCacheRecorder) GetOutboxWatermark(ctx context.Context) (int64, error) {
  107. return 0, nil
  108. }
  109. func (s *schedulerCacheRecorder) SetOutboxWatermark(ctx context.Context, id int64) error {
  110. return nil
  111. }
  112. func (s *AccountRepoSuite) SetupTest() {
  113. s.ctx = context.Background()
  114. tx := testEntTx(s.T())
  115. s.client = tx.Client()
  116. s.repo = newAccountRepositoryWithSQL(s.client, tx, nil)
  117. }
  118. func TestAccountRepoSuite(t *testing.T) {
  119. suite.Run(t, new(AccountRepoSuite))
  120. }
  121. // --- Create / GetByID / Update / Delete ---
  122. func (s *AccountRepoSuite) TestCreate() {
  123. account := &service.Account{
  124. Name: "test-create",
  125. Platform: service.PlatformAnthropic,
  126. Type: service.AccountTypeOAuth,
  127. Status: service.StatusActive,
  128. Credentials: map[string]any{},
  129. Extra: map[string]any{},
  130. Concurrency: 3,
  131. Priority: 50,
  132. Schedulable: true,
  133. }
  134. err := s.repo.Create(s.ctx, account)
  135. s.Require().NoError(err, "Create")
  136. s.Require().NotZero(account.ID, "expected ID to be set")
  137. got, err := s.repo.GetByID(s.ctx, account.ID)
  138. s.Require().NoError(err, "GetByID")
  139. s.Require().Equal("test-create", got.Name)
  140. }
  141. func (s *AccountRepoSuite) TestGetByID_NotFound() {
  142. _, err := s.repo.GetByID(s.ctx, 999999)
  143. s.Require().Error(err, "expected error for non-existent ID")
  144. }
  145. func (s *AccountRepoSuite) TestUpdate() {
  146. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "original"})
  147. account.Name = "updated"
  148. err := s.repo.Update(s.ctx, account)
  149. s.Require().NoError(err, "Update")
  150. got, err := s.repo.GetByID(s.ctx, account.ID)
  151. s.Require().NoError(err, "GetByID after update")
  152. s.Require().Equal("updated", got.Name)
  153. }
  154. func (s *AccountRepoSuite) TestUpdate_SyncSchedulerSnapshotOnDisabled() {
  155. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "sync-update", Status: service.StatusActive, Schedulable: true})
  156. cacheRecorder := &schedulerCacheRecorder{}
  157. s.repo.schedulerCache = cacheRecorder
  158. account.Status = service.StatusDisabled
  159. err := s.repo.Update(s.ctx, account)
  160. s.Require().NoError(err, "Update")
  161. s.Require().Len(cacheRecorder.setAccounts, 1)
  162. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  163. s.Require().Equal(service.StatusDisabled, cacheRecorder.setAccounts[0].Status)
  164. }
  165. func (s *AccountRepoSuite) TestUpdate_SyncSchedulerSnapshotOnCredentialsChange() {
  166. account := mustCreateAccount(s.T(), s.client, &service.Account{
  167. Name: "sync-credentials-update",
  168. Status: service.StatusActive,
  169. Schedulable: true,
  170. Credentials: map[string]any{
  171. "model_mapping": map[string]any{
  172. "gpt-5": "gpt-5.1",
  173. },
  174. },
  175. })
  176. cacheRecorder := &schedulerCacheRecorder{}
  177. s.repo.schedulerCache = cacheRecorder
  178. account.Credentials = map[string]any{
  179. "model_mapping": map[string]any{
  180. "gpt-5": "gpt-5.2",
  181. },
  182. }
  183. err := s.repo.Update(s.ctx, account)
  184. s.Require().NoError(err, "Update")
  185. s.Require().Len(cacheRecorder.setAccounts, 1)
  186. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  187. mapping, ok := cacheRecorder.setAccounts[0].Credentials["model_mapping"].(map[string]any)
  188. s.Require().True(ok)
  189. s.Require().Equal("gpt-5.2", mapping["gpt-5"])
  190. }
  191. func (s *AccountRepoSuite) TestUpdateCredentials_SyncsSnapshotAndDurableOutbox() {
  192. account := mustCreateAccount(s.T(), s.client, &service.Account{
  193. Name: "sync-refresh-credentials",
  194. Status: service.StatusActive,
  195. Schedulable: true,
  196. Credentials: map[string]any{"access_token": "old-token"},
  197. })
  198. cacheRecorder := &schedulerCacheRecorder{}
  199. s.repo.schedulerCache = cacheRecorder
  200. _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  201. s.Require().NoError(err)
  202. s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{"access_token": "new-token"}))
  203. s.Require().Len(cacheRecorder.setAccounts, 1)
  204. s.Require().Equal("new-token", cacheRecorder.setAccounts[0].GetCredential("access_token"))
  205. var outboxCount int
  206. err = scanSingleRow(
  207. s.ctx,
  208. s.repo.sql,
  209. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  210. []any{service.SchedulerOutboxEventAccountChanged, account.ID},
  211. &outboxCount,
  212. )
  213. s.Require().NoError(err)
  214. s.Require().Equal(1, outboxCount)
  215. }
  216. func (s *AccountRepoSuite) TestDelete() {
  217. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "to-delete"})
  218. err := s.repo.Delete(s.ctx, account.ID)
  219. s.Require().NoError(err, "Delete")
  220. _, err = s.repo.GetByID(s.ctx, account.ID)
  221. s.Require().Error(err, "expected error after delete")
  222. }
  223. func (s *AccountRepoSuite) TestDelete_RemovesSchedulerAccountSnapshot() {
  224. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "to-delete-cache"})
  225. cacheRecorder := &schedulerCacheRecorder{
  226. accounts: map[int64]*service.Account{
  227. account.ID: {
  228. ID: account.ID,
  229. Name: account.Name,
  230. Status: service.StatusActive,
  231. Schedulable: true,
  232. },
  233. },
  234. }
  235. s.repo.schedulerCache = cacheRecorder
  236. err := s.repo.Delete(s.ctx, account.ID)
  237. s.Require().NoError(err, "Delete")
  238. s.Require().Equal([]int64{account.ID}, cacheRecorder.deleteIDs)
  239. s.Require().NotContains(cacheRecorder.accounts, account.ID)
  240. }
  241. func (s *AccountRepoSuite) TestDelete_WithGroupBindings() {
  242. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-del"})
  243. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-del"})
  244. mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
  245. err := s.repo.Delete(s.ctx, account.ID)
  246. s.Require().NoError(err, "Delete should cascade remove bindings")
  247. count, err := s.client.AccountGroup.Query().Where(accountgroup.AccountIDEQ(account.ID)).Count(s.ctx)
  248. s.Require().NoError(err)
  249. s.Require().Zero(count, "expected bindings to be removed")
  250. }
  251. // --- List / ListWithFilters ---
  252. func (s *AccountRepoSuite) TestList() {
  253. mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc1"})
  254. mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc2"})
  255. accounts, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10})
  256. s.Require().NoError(err, "List")
  257. s.Require().Len(accounts, 2)
  258. s.Require().Equal(int64(2), page.Total)
  259. }
  260. func (s *AccountRepoSuite) TestListOAuthRefreshCandidatePage_GrokCursorAndExclusions() {
  261. now := time.Now().UTC()
  262. valid1 := mustCreateAccount(s.T(), s.client, &service.Account{
  263. Name: "grok-oauth-page-1",
  264. Platform: service.PlatformGrok,
  265. Type: service.AccountTypeOAuth,
  266. Status: service.StatusActive,
  267. Credentials: map[string]any{
  268. "access_token": "access-1",
  269. "refresh_token": "refresh-1",
  270. "expires_at": now.Add(30 * time.Minute).Format(time.RFC3339),
  271. },
  272. })
  273. unschedulable := mustCreateAccount(s.T(), s.client, &service.Account{
  274. Name: "grok-oauth-unschedulable-excluded",
  275. Platform: service.PlatformGrok,
  276. Type: service.AccountTypeOAuth,
  277. Status: service.StatusActive,
  278. Credentials: map[string]any{"refresh_token": "refresh-unschedulable"},
  279. })
  280. s.Require().NoError(s.client.Account.UpdateOneID(unschedulable.ID).SetSchedulable(false).Exec(s.ctx))
  281. mustCreateAccount(s.T(), s.client, &service.Account{
  282. Name: "grok-api-key-excluded",
  283. Platform: service.PlatformGrok,
  284. Type: service.AccountTypeAPIKey,
  285. Status: service.StatusActive,
  286. Credentials: map[string]any{
  287. "api_key": "api-key",
  288. "refresh_token": "must-not-make-api-key-eligible",
  289. },
  290. })
  291. valid2 := mustCreateAccount(s.T(), s.client, &service.Account{
  292. Name: "grok-oauth-page-2",
  293. Platform: service.PlatformGrok,
  294. Type: service.AccountTypeOAuth,
  295. Status: service.StatusActive,
  296. Credentials: map[string]any{"refresh_token": "refresh-2"},
  297. })
  298. mustCreateAccount(s.T(), s.client, &service.Account{
  299. Name: "grok-oauth-blank-refresh-excluded",
  300. Platform: service.PlatformGrok,
  301. Type: service.AccountTypeOAuth,
  302. Status: service.StatusActive,
  303. Credentials: map[string]any{"refresh_token": " "},
  304. })
  305. valid3 := mustCreateAccount(s.T(), s.client, &service.Account{
  306. Name: "grok-oauth-page-3",
  307. Platform: service.PlatformGrok,
  308. Type: service.AccountTypeOAuth,
  309. Status: service.StatusActive,
  310. Credentials: map[string]any{"refresh_token": "refresh-3"},
  311. })
  312. cooldown := mustCreateAccount(s.T(), s.client, &service.Account{
  313. Name: "grok-oauth-retry-cooldown-excluded",
  314. Platform: service.PlatformGrok,
  315. Type: service.AccountTypeOAuth,
  316. Status: service.StatusActive,
  317. Credentials: map[string]any{"refresh_token": "refresh-cooldown"},
  318. })
  319. s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, cooldown.ID, now.Add(10*time.Minute), "token refresh retry exhausted: timeout"))
  320. mustCreateAccount(s.T(), s.client, &service.Account{
  321. Name: "openai-oauth-excluded",
  322. Platform: service.PlatformOpenAI,
  323. Type: service.AccountTypeOAuth,
  324. Status: service.StatusActive,
  325. Credentials: map[string]any{"refresh_token": "refresh-openai"},
  326. })
  327. options := service.OAuthRefreshPageOptions{
  328. Platforms: []string{service.PlatformGrok},
  329. Limit: 2,
  330. ActiveOnly: true,
  331. RequireRefreshToken: true,
  332. ExcludeRetryCooldown: true,
  333. }
  334. firstPage, err := s.repo.ListOAuthRefreshCandidatePage(s.ctx, options)
  335. s.Require().NoError(err)
  336. first := firstPage.Accounts
  337. s.Require().Len(first, 2)
  338. s.Require().Equal([]int64{valid1.ID, valid2.ID}, []int64{first[0].ID, first[1].ID})
  339. s.Require().NotContains([]int64{first[0].ID, first[1].ID}, unschedulable.ID)
  340. options.AfterID = first[len(first)-1].ID
  341. secondPage, err := s.repo.ListOAuthRefreshCandidatePage(s.ctx, options)
  342. s.Require().NoError(err)
  343. second := secondPage.Accounts
  344. s.Require().Len(second, 1)
  345. s.Require().Equal(valid3.ID, second[0].ID)
  346. s.Require().NotContains([]int64{first[0].ID, first[1].ID}, second[0].ID)
  347. }
  348. func (s *AccountRepoSuite) TestListWithFilters() {
  349. tests := []struct {
  350. name string
  351. setup func(client *dbent.Client)
  352. platform string
  353. accType string
  354. status string
  355. search string
  356. groupID int64
  357. privacyMode string
  358. wantCount int
  359. validate func(accounts []service.Account)
  360. }{
  361. {
  362. name: "filter_by_platform",
  363. setup: func(client *dbent.Client) {
  364. mustCreateAccount(s.T(), client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic})
  365. mustCreateAccount(s.T(), client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI})
  366. },
  367. platform: service.PlatformOpenAI,
  368. wantCount: 1,
  369. validate: func(accounts []service.Account) {
  370. s.Require().Equal(service.PlatformOpenAI, accounts[0].Platform)
  371. },
  372. },
  373. {
  374. name: "filter_by_type",
  375. setup: func(client *dbent.Client) {
  376. mustCreateAccount(s.T(), client, &service.Account{Name: "t1", Type: service.AccountTypeOAuth})
  377. mustCreateAccount(s.T(), client, &service.Account{Name: "t2", Type: service.AccountTypeAPIKey})
  378. },
  379. accType: service.AccountTypeAPIKey,
  380. wantCount: 1,
  381. validate: func(accounts []service.Account) {
  382. s.Require().Equal(service.AccountTypeAPIKey, accounts[0].Type)
  383. },
  384. },
  385. {
  386. name: "filter_by_status",
  387. setup: func(client *dbent.Client) {
  388. mustCreateAccount(s.T(), client, &service.Account{Name: "s1", Status: service.StatusActive})
  389. mustCreateAccount(s.T(), client, &service.Account{Name: "s2", Status: service.StatusDisabled})
  390. },
  391. status: service.StatusDisabled,
  392. wantCount: 1,
  393. validate: func(accounts []service.Account) {
  394. s.Require().Equal(service.StatusDisabled, accounts[0].Status)
  395. },
  396. },
  397. {
  398. name: "filter_by_status_active_excludes_runtime_blocked_accounts",
  399. setup: func(client *dbent.Client) {
  400. mustCreateAccount(s.T(), client, &service.Account{Name: "active-normal", Status: service.StatusActive})
  401. rateLimited := mustCreateAccount(s.T(), client, &service.Account{Name: "active-rate-limited", Status: service.StatusActive})
  402. err := client.Account.UpdateOneID(rateLimited.ID).
  403. SetRateLimitResetAt(time.Now().Add(10 * time.Minute)).
  404. Exec(context.Background())
  405. s.Require().NoError(err)
  406. tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive})
  407. err = client.Account.UpdateOneID(tempUnsched.ID).
  408. SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
  409. Exec(context.Background())
  410. s.Require().NoError(err)
  411. unsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-unsched", Status: service.StatusActive})
  412. err = client.Account.UpdateOneID(unsched.ID).
  413. SetSchedulable(false).
  414. Exec(context.Background())
  415. s.Require().NoError(err)
  416. },
  417. status: service.StatusActive,
  418. wantCount: 1,
  419. validate: func(accounts []service.Account) {
  420. s.Require().Equal("active-normal", accounts[0].Name)
  421. },
  422. },
  423. {
  424. name: "filter_by_status_unschedulable_excludes_rate_limited_and_temp_unschedulable",
  425. setup: func(client *dbent.Client) {
  426. mustCreateAccount(s.T(), client, &service.Account{Name: "active-normal", Status: service.StatusActive, Schedulable: true})
  427. unsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-unsched", Status: service.StatusActive})
  428. err := client.Account.UpdateOneID(unsched.ID).
  429. SetSchedulable(false).
  430. Exec(context.Background())
  431. s.Require().NoError(err)
  432. rateLimited := mustCreateAccount(s.T(), client, &service.Account{Name: "active-rate-limited", Status: service.StatusActive})
  433. err = client.Account.UpdateOneID(rateLimited.ID).
  434. SetSchedulable(false).
  435. SetRateLimitResetAt(time.Now().Add(10 * time.Minute)).
  436. Exec(context.Background())
  437. s.Require().NoError(err)
  438. tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive})
  439. err = client.Account.UpdateOneID(tempUnsched.ID).
  440. SetSchedulable(false).
  441. SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
  442. Exec(context.Background())
  443. s.Require().NoError(err)
  444. },
  445. status: "unschedulable",
  446. wantCount: 1,
  447. validate: func(accounts []service.Account) {
  448. s.Require().Equal("active-unsched", accounts[0].Name)
  449. },
  450. },
  451. {
  452. name: "filter_by_status_rate_limited_excludes_temp_unschedulable",
  453. setup: func(client *dbent.Client) {
  454. rateLimited := mustCreateAccount(s.T(), client, &service.Account{Name: "active-rate-limited", Status: service.StatusActive})
  455. err := client.Account.UpdateOneID(rateLimited.ID).
  456. SetRateLimitResetAt(time.Now().Add(10 * time.Minute)).
  457. Exec(context.Background())
  458. s.Require().NoError(err)
  459. tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive})
  460. err = client.Account.UpdateOneID(tempUnsched.ID).
  461. SetRateLimitResetAt(time.Now().Add(20 * time.Minute)).
  462. SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
  463. Exec(context.Background())
  464. s.Require().NoError(err)
  465. },
  466. status: "rate_limited",
  467. wantCount: 1,
  468. validate: func(accounts []service.Account) {
  469. s.Require().Equal("active-rate-limited", accounts[0].Name)
  470. },
  471. },
  472. {
  473. name: "filter_by_status_temp_unschedulable_excludes_manually_unschedulable",
  474. setup: func(client *dbent.Client) {
  475. tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive, Schedulable: true})
  476. err := client.Account.UpdateOneID(tempUnsched.ID).
  477. SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
  478. Exec(context.Background())
  479. s.Require().NoError(err)
  480. unsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-unsched", Status: service.StatusActive})
  481. err = client.Account.UpdateOneID(unsched.ID).
  482. SetSchedulable(false).
  483. Exec(context.Background())
  484. s.Require().NoError(err)
  485. },
  486. status: "temp_unschedulable",
  487. wantCount: 1,
  488. validate: func(accounts []service.Account) {
  489. s.Require().Equal("active-temp-unsched", accounts[0].Name)
  490. },
  491. },
  492. {
  493. name: "filter_by_search",
  494. setup: func(client *dbent.Client) {
  495. mustCreateAccount(s.T(), client, &service.Account{Name: "alpha-account"})
  496. mustCreateAccount(s.T(), client, &service.Account{Name: "beta-account"})
  497. },
  498. search: "alpha",
  499. wantCount: 1,
  500. validate: func(accounts []service.Account) {
  501. s.Require().Contains(accounts[0].Name, "alpha")
  502. },
  503. },
  504. {
  505. name: "filter_by_ungrouped",
  506. setup: func(client *dbent.Client) {
  507. group := mustCreateGroup(s.T(), client, &service.Group{Name: "g-ungrouped"})
  508. grouped := mustCreateAccount(s.T(), client, &service.Account{Name: "grouped-account"})
  509. mustCreateAccount(s.T(), client, &service.Account{Name: "ungrouped-account"})
  510. mustBindAccountToGroup(s.T(), client, grouped.ID, group.ID, 1)
  511. },
  512. groupID: service.AccountListGroupUngrouped,
  513. wantCount: 1,
  514. validate: func(accounts []service.Account) {
  515. s.Require().Equal("ungrouped-account", accounts[0].Name)
  516. s.Require().Empty(accounts[0].GroupIDs)
  517. },
  518. },
  519. {
  520. name: "filter_by_privacy_mode",
  521. setup: func(client *dbent.Client) {
  522. mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-ok", Extra: map[string]any{"privacy_mode": service.PrivacyModeTrainingOff}})
  523. mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-fail", Extra: map[string]any{"privacy_mode": service.PrivacyModeFailed}})
  524. },
  525. privacyMode: service.PrivacyModeTrainingOff,
  526. wantCount: 1,
  527. validate: func(accounts []service.Account) {
  528. s.Require().Equal("privacy-ok", accounts[0].Name)
  529. },
  530. },
  531. {
  532. name: "filter_by_privacy_mode_unset",
  533. setup: func(client *dbent.Client) {
  534. mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-unset", Extra: nil})
  535. mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-empty", Extra: map[string]any{"privacy_mode": ""}})
  536. mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-set", Extra: map[string]any{"privacy_mode": service.PrivacyModeTrainingOff}})
  537. },
  538. privacyMode: service.AccountPrivacyModeUnsetFilter,
  539. wantCount: 2,
  540. validate: func(accounts []service.Account) {
  541. names := []string{accounts[0].Name, accounts[1].Name}
  542. s.ElementsMatch([]string{"privacy-unset", "privacy-empty"}, names)
  543. },
  544. },
  545. }
  546. for _, tt := range tests {
  547. s.Run(tt.name, func() {
  548. // 每个 case 重新获取隔离资源
  549. tx := testEntTx(s.T())
  550. client := tx.Client()
  551. repo := newAccountRepositoryWithSQL(client, tx, nil)
  552. ctx := context.Background()
  553. tt.setup(client)
  554. accounts, page, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, tt.platform, tt.accType, tt.status, tt.search, tt.groupID, tt.privacyMode)
  555. s.Require().NoError(err)
  556. s.Require().Len(accounts, tt.wantCount)
  557. // Regression guard for issue #3601: when the whole result set fits on a single page,
  558. // pagination.Total must match len(items). A mismatch means the Count query was applied
  559. // against different predicates than the list query — the exact symptom reported.
  560. s.Require().NotNil(page)
  561. s.Require().Equal(int64(tt.wantCount), page.Total, "total must match items on single page")
  562. if tt.validate != nil {
  563. tt.validate(accounts)
  564. }
  565. })
  566. }
  567. }
  568. // --- ListByGroup / ListActive / ListByPlatform ---
  569. func (s *AccountRepoSuite) TestListByGroup() {
  570. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-list"})
  571. acc1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Status: service.StatusActive})
  572. acc2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Status: service.StatusActive})
  573. mustBindAccountToGroup(s.T(), s.client, acc1.ID, group.ID, 2)
  574. mustBindAccountToGroup(s.T(), s.client, acc2.ID, group.ID, 1)
  575. accounts, err := s.repo.ListByGroup(s.ctx, group.ID)
  576. s.Require().NoError(err, "ListByGroup")
  577. s.Require().Len(accounts, 2)
  578. // Should be ordered by priority
  579. s.Require().Equal(acc2.ID, accounts[0].ID, "expected acc2 first (priority=1)")
  580. }
  581. func (s *AccountRepoSuite) TestListActive() {
  582. mustCreateAccount(s.T(), s.client, &service.Account{Name: "active1", Status: service.StatusActive})
  583. mustCreateAccount(s.T(), s.client, &service.Account{Name: "inactive1", Status: service.StatusDisabled})
  584. accounts, err := s.repo.ListActive(s.ctx)
  585. s.Require().NoError(err, "ListActive")
  586. s.Require().Len(accounts, 1)
  587. s.Require().Equal("active1", accounts[0].Name)
  588. }
  589. func (s *AccountRepoSuite) TestListByPlatform() {
  590. mustCreateAccount(s.T(), s.client, &service.Account{Name: "p1", Platform: service.PlatformAnthropic, Status: service.StatusActive})
  591. mustCreateAccount(s.T(), s.client, &service.Account{Name: "p2", Platform: service.PlatformOpenAI, Status: service.StatusActive})
  592. accounts, err := s.repo.ListByPlatform(s.ctx, service.PlatformAnthropic)
  593. s.Require().NoError(err, "ListByPlatform")
  594. s.Require().Len(accounts, 1)
  595. s.Require().Equal(service.PlatformAnthropic, accounts[0].Platform)
  596. }
  597. // --- Preload and VirtualFields ---
  598. func (s *AccountRepoSuite) TestPreload_And_VirtualFields() {
  599. proxy := mustCreateProxy(s.T(), s.client, &service.Proxy{Name: "p1"})
  600. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g1"})
  601. account := mustCreateAccount(s.T(), s.client, &service.Account{
  602. Name: "acc1",
  603. ProxyID: &proxy.ID,
  604. })
  605. mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
  606. got, err := s.repo.GetByID(s.ctx, account.ID)
  607. s.Require().NoError(err, "GetByID")
  608. s.Require().NotNil(got.Proxy, "expected Proxy preload")
  609. s.Require().Equal(proxy.ID, got.Proxy.ID)
  610. s.Require().Len(got.GroupIDs, 1, "expected GroupIDs to be populated")
  611. s.Require().Equal(group.ID, got.GroupIDs[0])
  612. s.Require().Len(got.Groups, 1, "expected Groups to be populated")
  613. s.Require().Equal(group.ID, got.Groups[0].ID)
  614. accounts, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, "", "", "", "acc", 0, "")
  615. s.Require().NoError(err, "ListWithFilters")
  616. s.Require().Equal(int64(1), page.Total)
  617. s.Require().Len(accounts, 1)
  618. s.Require().NotNil(accounts[0].Proxy, "expected Proxy preload in list")
  619. s.Require().Equal(proxy.ID, accounts[0].Proxy.ID)
  620. s.Require().Len(accounts[0].GroupIDs, 1, "expected GroupIDs in list")
  621. s.Require().Equal(group.ID, accounts[0].GroupIDs[0])
  622. }
  623. // --- GroupBinding / AddToGroup / RemoveFromGroup / BindGroups / GetGroups ---
  624. func (s *AccountRepoSuite) TestGroupBinding_And_BindGroups() {
  625. g1 := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g1"})
  626. g2 := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g2"})
  627. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc"})
  628. s.Require().NoError(s.repo.AddToGroup(s.ctx, account.ID, g1.ID, 10), "AddToGroup")
  629. groups, err := s.repo.GetGroups(s.ctx, account.ID)
  630. s.Require().NoError(err, "GetGroups")
  631. s.Require().Len(groups, 1, "expected 1 group")
  632. s.Require().Equal(g1.ID, groups[0].ID)
  633. s.Require().NoError(s.repo.RemoveFromGroup(s.ctx, account.ID, g1.ID), "RemoveFromGroup")
  634. groups, err = s.repo.GetGroups(s.ctx, account.ID)
  635. s.Require().NoError(err, "GetGroups after remove")
  636. s.Require().Empty(groups, "expected 0 groups after remove")
  637. s.Require().NoError(s.repo.BindGroups(s.ctx, account.ID, []int64{g1.ID, g2.ID}), "BindGroups")
  638. groups, err = s.repo.GetGroups(s.ctx, account.ID)
  639. s.Require().NoError(err, "GetGroups after bind")
  640. s.Require().Len(groups, 2, "expected 2 groups after bind")
  641. }
  642. func (s *AccountRepoSuite) TestBindGroups_EmptyList() {
  643. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-empty"})
  644. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-empty"})
  645. mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
  646. s.Require().NoError(s.repo.BindGroups(s.ctx, account.ID, []int64{}), "BindGroups empty")
  647. groups, err := s.repo.GetGroups(s.ctx, account.ID)
  648. s.Require().NoError(err)
  649. s.Require().Empty(groups, "expected 0 groups after binding empty list")
  650. }
  651. // --- Schedulable ---
  652. func (s *AccountRepoSuite) TestListSchedulable() {
  653. now := time.Now()
  654. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sched"})
  655. okAcc := mustCreateAccount(s.T(), s.client, &service.Account{Name: "ok", Schedulable: true})
  656. mustBindAccountToGroup(s.T(), s.client, okAcc.ID, group.ID, 1)
  657. future := now.Add(10 * time.Minute)
  658. overloaded := mustCreateAccount(s.T(), s.client, &service.Account{Name: "over", Schedulable: true, OverloadUntil: &future})
  659. mustBindAccountToGroup(s.T(), s.client, overloaded.ID, group.ID, 1)
  660. sched, err := s.repo.ListSchedulable(s.ctx)
  661. s.Require().NoError(err, "ListSchedulable")
  662. ids := idsOfAccounts(sched)
  663. s.Require().Contains(ids, okAcc.ID)
  664. s.Require().NotContains(ids, overloaded.ID)
  665. }
  666. func (s *AccountRepoSuite) TestListSchedulableByGroupID_TimeBoundaries_And_StatusUpdates() {
  667. now := time.Now()
  668. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sched"})
  669. okAcc := mustCreateAccount(s.T(), s.client, &service.Account{Name: "ok", Schedulable: true})
  670. mustBindAccountToGroup(s.T(), s.client, okAcc.ID, group.ID, 1)
  671. future := now.Add(10 * time.Minute)
  672. overloaded := mustCreateAccount(s.T(), s.client, &service.Account{Name: "over", Schedulable: true, OverloadUntil: &future})
  673. mustBindAccountToGroup(s.T(), s.client, overloaded.ID, group.ID, 1)
  674. rateLimited := mustCreateAccount(s.T(), s.client, &service.Account{Name: "rl", Schedulable: true})
  675. mustBindAccountToGroup(s.T(), s.client, rateLimited.ID, group.ID, 1)
  676. s.Require().NoError(s.repo.SetRateLimited(s.ctx, rateLimited.ID, now.Add(10*time.Minute)), "SetRateLimited")
  677. s.Require().NoError(s.repo.SetError(s.ctx, overloaded.ID, "boom"), "SetError")
  678. sched, err := s.repo.ListSchedulableByGroupID(s.ctx, group.ID)
  679. s.Require().NoError(err, "ListSchedulableByGroupID")
  680. s.Require().Len(sched, 1, "expected only ok account schedulable")
  681. s.Require().Equal(okAcc.ID, sched[0].ID)
  682. s.Require().NoError(s.repo.ClearRateLimit(s.ctx, rateLimited.ID), "ClearRateLimit")
  683. sched2, err := s.repo.ListSchedulableByGroupID(s.ctx, group.ID)
  684. s.Require().NoError(err, "ListSchedulableByGroupID after ClearRateLimit")
  685. s.Require().Len(sched2, 2, "expected 2 schedulable accounts after ClearRateLimit")
  686. }
  687. func (s *AccountRepoSuite) TestListSchedulableByPlatform() {
  688. mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic, Schedulable: true})
  689. mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI, Schedulable: true})
  690. accounts, err := s.repo.ListSchedulableByPlatform(s.ctx, service.PlatformAnthropic)
  691. s.Require().NoError(err)
  692. s.Require().Len(accounts, 1)
  693. s.Require().Equal(service.PlatformAnthropic, accounts[0].Platform)
  694. }
  695. func (s *AccountRepoSuite) TestListSchedulableByGroupIDAndPlatform() {
  696. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sp"})
  697. a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic, Schedulable: true})
  698. a2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI, Schedulable: true})
  699. mustBindAccountToGroup(s.T(), s.client, a1.ID, group.ID, 1)
  700. mustBindAccountToGroup(s.T(), s.client, a2.ID, group.ID, 2)
  701. accounts, err := s.repo.ListSchedulableByGroupIDAndPlatform(s.ctx, group.ID, service.PlatformAnthropic)
  702. s.Require().NoError(err)
  703. s.Require().Len(accounts, 1)
  704. s.Require().Equal(a1.ID, accounts[0].ID)
  705. }
  706. func (s *AccountRepoSuite) TestSetSchedulable() {
  707. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-sched", Schedulable: true})
  708. cacheRecorder := &schedulerCacheRecorder{}
  709. s.repo.schedulerCache = cacheRecorder
  710. s.Require().NoError(s.repo.SetSchedulable(s.ctx, account.ID, false))
  711. got, err := s.repo.GetByID(s.ctx, account.ID)
  712. s.Require().NoError(err)
  713. s.Require().False(got.Schedulable)
  714. s.Require().Len(cacheRecorder.setAccounts, 1)
  715. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  716. }
  717. func (s *AccountRepoSuite) TestBulkUpdate_SyncSchedulerSnapshotOnDisabled() {
  718. account1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-1", Status: service.StatusActive, Schedulable: true})
  719. account2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-2", Status: service.StatusActive, Schedulable: true})
  720. cacheRecorder := &schedulerCacheRecorder{}
  721. s.repo.schedulerCache = cacheRecorder
  722. disabled := service.StatusDisabled
  723. rows, err := s.repo.BulkUpdate(s.ctx, []int64{account1.ID, account2.ID}, service.AccountBulkUpdate{
  724. Status: &disabled,
  725. })
  726. s.Require().NoError(err)
  727. s.Require().Equal(int64(2), rows)
  728. s.Require().Len(cacheRecorder.setAccounts, 2)
  729. ids := map[int64]struct{}{}
  730. for _, acc := range cacheRecorder.setAccounts {
  731. ids[acc.ID] = struct{}{}
  732. }
  733. s.Require().Contains(ids, account1.ID)
  734. s.Require().Contains(ids, account2.ID)
  735. }
  736. // --- SetOverloaded / SetRateLimited / ClearRateLimit ---
  737. func (s *AccountRepoSuite) TestSetOverloaded() {
  738. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-over"})
  739. until := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
  740. cacheRecorder := &schedulerCacheRecorder{}
  741. s.repo.schedulerCache = cacheRecorder
  742. s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
  743. got, err := s.repo.GetByID(s.ctx, account.ID)
  744. s.Require().NoError(err)
  745. s.Require().NotNil(got.OverloadUntil)
  746. s.Require().WithinDuration(until, *got.OverloadUntil, time.Second)
  747. s.Require().Len(cacheRecorder.setAccounts, 1)
  748. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  749. s.Require().NotNil(cacheRecorder.setAccounts[0].OverloadUntil)
  750. s.Require().WithinDuration(until, *cacheRecorder.setAccounts[0].OverloadUntil, time.Second)
  751. }
  752. func (s *AccountRepoSuite) TestSetRateLimited() {
  753. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl"})
  754. resetAt := time.Date(2025, 6, 15, 14, 0, 0, 0, time.UTC)
  755. s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, resetAt))
  756. got, err := s.repo.GetByID(s.ctx, account.ID)
  757. s.Require().NoError(err)
  758. s.Require().NotNil(got.RateLimitedAt)
  759. s.Require().NotNil(got.RateLimitResetAt)
  760. s.Require().WithinDuration(resetAt, *got.RateLimitResetAt, time.Second)
  761. }
  762. func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() {
  763. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl-monotonic"})
  764. later := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
  765. earlier := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second)
  766. cacheRecorder := &schedulerCacheRecorder{}
  767. s.repo.schedulerCache = cacheRecorder
  768. s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, later))
  769. s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, earlier))
  770. got, err := s.repo.GetByID(s.ctx, account.ID)
  771. s.Require().NoError(err)
  772. s.Require().NotNil(got.RateLimitResetAt)
  773. s.Require().WithinDuration(later, *got.RateLimitResetAt, time.Second)
  774. s.Require().Len(cacheRecorder.setAccounts, 2)
  775. s.Require().NotNil(cacheRecorder.setAccounts[1].RateLimitResetAt)
  776. s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second)
  777. }
  778. func (s *AccountRepoSuite) TestClearRateLimitIfObservedProtectsRearmed429Generation() {
  779. account := mustCreateAccount(s.T(), s.client, &service.Account{
  780. Name: "acc-rl-conditional-clear",
  781. Platform: service.PlatformGrok,
  782. Type: service.AccountTypeOAuth,
  783. })
  784. firstReset := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
  785. rearmedReset := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second)
  786. s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, firstReset))
  787. staleGeneration, err := s.repo.GetByID(s.ctx, account.ID)
  788. s.Require().NoError(err)
  789. s.Require().NotNil(staleGeneration.RateLimitedAt)
  790. s.Require().NotNil(staleGeneration.RateLimitResetAt)
  791. cleared, err := s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt)
  792. s.Require().NoError(err)
  793. s.Require().True(cleared)
  794. // A newer generation may legitimately re-arm a shorter boundary after the
  795. // first generation was cleared. The stale success must not erase it.
  796. s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, rearmedReset))
  797. cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt)
  798. s.Require().NoError(err)
  799. s.Require().False(cleared)
  800. got, err := s.repo.GetByID(s.ctx, account.ID)
  801. s.Require().NoError(err)
  802. s.Require().NotNil(got.RateLimitedAt)
  803. s.Require().NotNil(got.RateLimitResetAt)
  804. s.Require().WithinDuration(rearmedReset, *got.RateLimitResetAt, time.Second)
  805. // An admin can retype the row while the successful OAuth request is still
  806. // in flight. The stale OAuth recovery must not cross into API-key state even
  807. // when both observed timestamps still match.
  808. _, err = s.client.Account.UpdateOneID(account.ID).
  809. SetType(service.AccountTypeAPIKey).
  810. Save(s.ctx)
  811. s.Require().NoError(err)
  812. cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *got.RateLimitedAt, *got.RateLimitResetAt)
  813. s.Require().NoError(err)
  814. s.Require().False(cleared)
  815. retyped, err := s.repo.GetByID(s.ctx, account.ID)
  816. s.Require().NoError(err)
  817. s.Require().Equal(service.AccountTypeAPIKey, retyped.Type)
  818. s.Require().NotNil(retyped.RateLimitedAt)
  819. s.Require().NotNil(retyped.RateLimitResetAt)
  820. s.Require().WithinDuration(rearmedReset, *retyped.RateLimitResetAt, time.Second)
  821. }
  822. func (s *AccountRepoSuite) TestClearRateLimit() {
  823. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"})
  824. until := time.Now().Add(1 * time.Hour)
  825. s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
  826. s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, until))
  827. s.Require().NoError(s.repo.ClearRateLimit(s.ctx, account.ID))
  828. got, err := s.repo.GetByID(s.ctx, account.ID)
  829. s.Require().NoError(err)
  830. s.Require().Nil(got.RateLimitedAt)
  831. s.Require().Nil(got.RateLimitResetAt)
  832. s.Require().Nil(got.OverloadUntil)
  833. }
  834. func (s *AccountRepoSuite) TestResetQuotaUsedAndClearRateLimitCooldownPreservesOtherRuntimeState() {
  835. account := mustCreateAccount(s.T(), s.client, &service.Account{
  836. Name: "acc-reset-quota-cooldown",
  837. Extra: map[string]any{
  838. "quota_used": 12.5,
  839. "quota_daily_used": 5.0,
  840. "quota_weekly_used": 9.0,
  841. "model_rate_limits": map[string]any{
  842. "claude-sonnet-4-5": map[string]any{"rate_limit_reset_at": "2026-09-01T10:00:00Z"},
  843. },
  844. },
  845. })
  846. until := time.Now().Add(1 * time.Hour)
  847. s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
  848. s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, until))
  849. s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, account.ID, until, "preserve-me"))
  850. cacheRecorder := &schedulerCacheRecorder{}
  851. s.repo.schedulerCache = cacheRecorder
  852. s.Require().NoError(s.repo.ResetQuotaUsedAndClearRateLimitCooldown(s.ctx, account.ID))
  853. got, err := s.repo.GetByID(s.ctx, account.ID)
  854. s.Require().NoError(err)
  855. s.Require().Nil(got.RateLimitedAt)
  856. s.Require().Nil(got.RateLimitResetAt)
  857. s.Require().NotNil(got.OverloadUntil)
  858. s.Require().WithinDuration(until, *got.OverloadUntil, time.Second)
  859. s.Require().NotNil(got.TempUnschedulableUntil)
  860. s.Require().WithinDuration(until, *got.TempUnschedulableUntil, time.Second)
  861. s.Require().Equal("preserve-me", got.TempUnschedulableReason)
  862. s.Require().Contains(got.Extra, "model_rate_limits")
  863. s.Require().Equal(float64(0), got.Extra["quota_used"])
  864. s.Require().Equal(float64(0), got.Extra["quota_daily_used"])
  865. s.Require().Equal(float64(0), got.Extra["quota_weekly_used"])
  866. var pendingEventExists bool
  867. s.Require().NoError(scanSingleRow(s.ctx, s.repo.sql, `
  868. SELECT EXISTS (
  869. SELECT 1 FROM scheduler_outbox
  870. WHERE event_type = $1 AND account_id = $2 AND dedup_key IS NOT NULL
  871. )`, []any{service.SchedulerOutboxEventAccountChanged, account.ID}, &pendingEventExists))
  872. s.Require().True(pendingEventExists)
  873. s.Require().Len(cacheRecorder.setAccounts, 1)
  874. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  875. }
  876. func (s *AccountRepoSuite) TestTempUnschedulableFieldsLoadedByGetByIDAndGetByIDs() {
  877. acc1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-temp-1"})
  878. acc2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-temp-2"})
  879. until := time.Now().Add(15 * time.Minute).UTC().Truncate(time.Second)
  880. reason := `{"rule":"429","matched_keyword":"too many requests"}`
  881. s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, acc1.ID, until, reason))
  882. gotByID, err := s.repo.GetByID(s.ctx, acc1.ID)
  883. s.Require().NoError(err)
  884. s.Require().NotNil(gotByID.TempUnschedulableUntil)
  885. s.Require().WithinDuration(until, *gotByID.TempUnschedulableUntil, time.Second)
  886. s.Require().Equal(reason, gotByID.TempUnschedulableReason)
  887. gotByIDs, err := s.repo.GetByIDs(s.ctx, []int64{acc2.ID, acc1.ID})
  888. s.Require().NoError(err)
  889. s.Require().Len(gotByIDs, 2)
  890. s.Require().Equal(acc2.ID, gotByIDs[0].ID)
  891. s.Require().Nil(gotByIDs[0].TempUnschedulableUntil)
  892. s.Require().Equal("", gotByIDs[0].TempUnschedulableReason)
  893. s.Require().Equal(acc1.ID, gotByIDs[1].ID)
  894. s.Require().NotNil(gotByIDs[1].TempUnschedulableUntil)
  895. s.Require().WithinDuration(until, *gotByIDs[1].TempUnschedulableUntil, time.Second)
  896. s.Require().Equal(reason, gotByIDs[1].TempUnschedulableReason)
  897. cacheRecorder := &schedulerCacheRecorder{}
  898. s.repo.schedulerCache = cacheRecorder
  899. s.Require().NoError(s.repo.ClearTempUnschedulable(s.ctx, acc1.ID))
  900. cleared, err := s.repo.GetByID(s.ctx, acc1.ID)
  901. s.Require().NoError(err)
  902. s.Require().Nil(cleared.TempUnschedulableUntil)
  903. s.Require().Equal("", cleared.TempUnschedulableReason)
  904. s.Require().Len(cacheRecorder.setAccounts, 1)
  905. s.Require().Equal(acc1.ID, cacheRecorder.setAccounts[0].ID)
  906. s.Require().Nil(cacheRecorder.setAccounts[0].TempUnschedulableUntil)
  907. s.Require().Equal("", cacheRecorder.setAccounts[0].TempUnschedulableReason)
  908. }
  909. func (s *AccountRepoSuite) TestSetTempUnschedulableSkipsOutboxWhenWindowDoesNotExtend() {
  910. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-temp-noop"})
  911. cacheRecorder := &schedulerCacheRecorder{}
  912. s.repo.schedulerCache = cacheRecorder
  913. _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  914. s.Require().NoError(err)
  915. until := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
  916. s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, account.ID, until, "first"))
  917. var count int
  918. err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
  919. s.Require().NoError(err)
  920. s.Require().Equal(1, count)
  921. s.Require().Len(cacheRecorder.setAccounts, 1)
  922. s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, account.ID, until.Add(-5*time.Minute), "older"))
  923. err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
  924. s.Require().NoError(err)
  925. s.Require().Equal(1, count)
  926. s.Require().Len(cacheRecorder.setAccounts, 1)
  927. got, err := s.repo.GetByID(s.ctx, account.ID)
  928. s.Require().NoError(err)
  929. s.Require().Equal("first", got.TempUnschedulableReason)
  930. s.Require().NotNil(got.TempUnschedulableUntil)
  931. s.Require().WithinDuration(until, *got.TempUnschedulableUntil, time.Second)
  932. }
  933. func (s *AccountRepoSuite) TestClearModelRateLimits_SyncsSchedulerSnapshot() {
  934. account := mustCreateAccount(s.T(), s.client, &service.Account{
  935. Name: "acc-clear-model-rate",
  936. Extra: map[string]any{
  937. "model_rate_limits": map[string]any{
  938. "claude-sonnet-4-5": map[string]any{
  939. "rate_limit_reset_at": "2026-06-03T10:00:00Z",
  940. },
  941. },
  942. },
  943. })
  944. cacheRecorder := &schedulerCacheRecorder{}
  945. s.repo.schedulerCache = cacheRecorder
  946. s.Require().NoError(s.repo.ClearModelRateLimits(s.ctx, account.ID))
  947. got, err := s.repo.GetByID(s.ctx, account.ID)
  948. s.Require().NoError(err)
  949. s.Require().NotContains(got.Extra, "model_rate_limits")
  950. s.Require().Len(cacheRecorder.setAccounts, 1)
  951. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  952. s.Require().NotContains(cacheRecorder.setAccounts[0].Extra, "model_rate_limits")
  953. }
  954. // --- UpdateLastUsed ---
  955. func (s *AccountRepoSuite) TestUpdateLastUsed() {
  956. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-used"})
  957. s.Require().Nil(account.LastUsedAt)
  958. s.Require().NoError(s.repo.UpdateLastUsed(s.ctx, account.ID))
  959. got, err := s.repo.GetByID(s.ctx, account.ID)
  960. s.Require().NoError(err)
  961. s.Require().NotNil(got.LastUsedAt)
  962. }
  963. // --- SetError ---
  964. func (s *AccountRepoSuite) TestSetError() {
  965. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-err", Status: service.StatusActive, Schedulable: true})
  966. cacheRecorder := &schedulerCacheRecorder{}
  967. s.repo.schedulerCache = cacheRecorder
  968. _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  969. s.Require().NoError(err)
  970. s.Require().NoError(s.repo.SetError(s.ctx, account.ID, "something went wrong"))
  971. got, err := s.repo.GetByID(s.ctx, account.ID)
  972. s.Require().NoError(err)
  973. s.Require().Equal(service.StatusError, got.Status)
  974. s.Require().Equal("something went wrong", got.ErrorMessage)
  975. s.Require().False(got.Schedulable)
  976. s.Require().Len(cacheRecorder.setAccounts, 1)
  977. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  978. s.Require().Equal(service.StatusError, cacheRecorder.setAccounts[0].Status)
  979. s.Require().False(cacheRecorder.setAccounts[0].Schedulable)
  980. var outboxCount int
  981. err = scanSingleRow(
  982. s.ctx,
  983. s.repo.sql,
  984. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  985. []any{service.SchedulerOutboxEventAccountChanged, account.ID},
  986. &outboxCount,
  987. )
  988. s.Require().NoError(err)
  989. s.Require().Equal(1, outboxCount)
  990. }
  991. func (s *AccountRepoSuite) TestSetGrokOAuthErrorIfCredentialsUnchanged_AppliesAndSyncsSchedulerState() {
  992. account := mustCreateAccount(s.T(), s.client, &service.Account{
  993. Name: "grok-conditional-error-applied",
  994. Platform: service.PlatformGrok,
  995. Type: service.AccountTypeOAuth,
  996. Status: service.StatusActive,
  997. Schedulable: true,
  998. Credentials: map[string]any{"access_token": "observed", "_token_version": int64(7)},
  999. })
  1000. observed, err := s.repo.GetByID(s.ctx, account.ID)
  1001. s.Require().NoError(err)
  1002. cacheRecorder := &schedulerCacheRecorder{}
  1003. s.repo.schedulerCache = cacheRecorder
  1004. _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  1005. s.Require().NoError(err)
  1006. applied, err := s.repo.SetGrokOAuthErrorIfCredentialsUnchanged(
  1007. s.ctx,
  1008. account.ID,
  1009. observed.Credentials,
  1010. "missing refresh token",
  1011. )
  1012. s.Require().NoError(err)
  1013. s.Require().True(applied)
  1014. got, err := s.repo.GetByID(s.ctx, account.ID)
  1015. s.Require().NoError(err)
  1016. s.Require().Equal(service.StatusError, got.Status)
  1017. s.Require().False(got.Schedulable)
  1018. s.Require().Equal("missing refresh token", got.ErrorMessage)
  1019. s.Require().Len(cacheRecorder.setAccounts, 1)
  1020. s.Require().Equal(service.StatusError, cacheRecorder.setAccounts[0].Status)
  1021. var outboxCount int
  1022. err = scanSingleRow(
  1023. s.ctx,
  1024. s.repo.sql,
  1025. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  1026. []any{service.SchedulerOutboxEventAccountChanged, account.ID},
  1027. &outboxCount,
  1028. )
  1029. s.Require().NoError(err)
  1030. s.Require().Equal(1, outboxCount)
  1031. }
  1032. func (s *AccountRepoSuite) TestSetGrokOAuthErrorIfCredentialsUnchanged_SkipsConcurrentReauthorization() {
  1033. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1034. Name: "grok-conditional-error-reauthorized",
  1035. Platform: service.PlatformGrok,
  1036. Type: service.AccountTypeOAuth,
  1037. Status: service.StatusActive,
  1038. Schedulable: true,
  1039. Credentials: map[string]any{"access_token": "observed", "_token_version": int64(7)},
  1040. })
  1041. observed, err := s.repo.GetByID(s.ctx, account.ID)
  1042. s.Require().NoError(err)
  1043. s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{
  1044. "access_token": "fresh-access",
  1045. "refresh_token": "fresh-refresh",
  1046. "expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
  1047. "_token_version": int64(8),
  1048. }))
  1049. cacheRecorder := &schedulerCacheRecorder{}
  1050. s.repo.schedulerCache = cacheRecorder
  1051. _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  1052. s.Require().NoError(err)
  1053. applied, err := s.repo.SetGrokOAuthErrorIfCredentialsUnchanged(
  1054. s.ctx,
  1055. account.ID,
  1056. observed.Credentials,
  1057. "stale reconciliation",
  1058. )
  1059. s.Require().NoError(err)
  1060. s.Require().False(applied)
  1061. got, err := s.repo.GetByID(s.ctx, account.ID)
  1062. s.Require().NoError(err)
  1063. s.Require().Equal(service.StatusActive, got.Status)
  1064. s.Require().True(got.Schedulable)
  1065. s.Require().Equal("fresh-refresh", got.GetGrokRefreshToken())
  1066. s.Require().Empty(cacheRecorder.setAccounts, "a lost compare-and-set race must not rewrite the scheduler snapshot")
  1067. var outboxCount int
  1068. err = scanSingleRow(
  1069. s.ctx,
  1070. s.repo.sql,
  1071. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  1072. []any{service.SchedulerOutboxEventAccountChanged, account.ID},
  1073. &outboxCount,
  1074. )
  1075. s.Require().NoError(err)
  1076. s.Require().Zero(outboxCount, "a lost compare-and-set race must not enqueue a stale account change")
  1077. }
  1078. func (s *AccountRepoSuite) TestUpdateGrokOAuthCredentialsIfUnchanged_AppliesAndPublishesSchedulerState() {
  1079. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1080. Name: "grok-refresh-success-cas-applied",
  1081. Platform: service.PlatformGrok,
  1082. Type: service.AccountTypeOAuth,
  1083. Status: service.StatusActive,
  1084. Schedulable: true,
  1085. Credentials: map[string]any{
  1086. "access_token": "attempted-access",
  1087. "refresh_token": "attempted-refresh",
  1088. "_token_version": int64(10),
  1089. },
  1090. })
  1091. observed, err := s.repo.GetByID(s.ctx, account.ID)
  1092. s.Require().NoError(err)
  1093. cacheRecorder := &schedulerCacheRecorder{}
  1094. s.repo.schedulerCache = cacheRecorder
  1095. _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  1096. s.Require().NoError(err)
  1097. applied, err := s.repo.UpdateGrokOAuthCredentialsIfUnchanged(
  1098. s.ctx,
  1099. account.ID,
  1100. observed.Credentials,
  1101. observed.ProxyID,
  1102. map[string]any{
  1103. "access_token": "rotated-access",
  1104. "refresh_token": "rotated-refresh",
  1105. "_token_version": int64(11),
  1106. },
  1107. )
  1108. s.Require().NoError(err)
  1109. s.Require().True(applied)
  1110. got, err := s.repo.GetByID(s.ctx, account.ID)
  1111. s.Require().NoError(err)
  1112. s.Require().Equal("rotated-refresh", got.GetGrokRefreshToken())
  1113. s.Require().Len(cacheRecorder.setAccounts, 1)
  1114. s.Require().Equal("rotated-refresh", cacheRecorder.setAccounts[0].GetGrokRefreshToken())
  1115. s.Require().NoError(cacheRecorder.setCtxErr)
  1116. var outboxCount int
  1117. err = scanSingleRow(
  1118. s.ctx,
  1119. s.repo.sql,
  1120. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  1121. []any{service.SchedulerOutboxEventAccountChanged, account.ID},
  1122. &outboxCount,
  1123. )
  1124. s.Require().NoError(err)
  1125. s.Require().Equal(1, outboxCount)
  1126. }
  1127. func (s *AccountRepoSuite) TestUpdateGrokOAuthCredentialsIfUnchanged_SkipsConcurrentReauthorization() {
  1128. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1129. Name: "grok-refresh-success-cas-reauthorized",
  1130. Platform: service.PlatformGrok,
  1131. Type: service.AccountTypeOAuth,
  1132. Status: service.StatusActive,
  1133. Schedulable: true,
  1134. Credentials: map[string]any{
  1135. "access_token": "attempted-access",
  1136. "refresh_token": "attempted-refresh",
  1137. "_token_version": int64(20),
  1138. },
  1139. })
  1140. observed, err := s.repo.GetByID(s.ctx, account.ID)
  1141. s.Require().NoError(err)
  1142. s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{
  1143. "access_token": "reauthorized-access",
  1144. "refresh_token": "reauthorized-refresh",
  1145. "_token_version": int64(21),
  1146. }))
  1147. cacheRecorder := &schedulerCacheRecorder{}
  1148. s.repo.schedulerCache = cacheRecorder
  1149. _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  1150. s.Require().NoError(err)
  1151. applied, err := s.repo.UpdateGrokOAuthCredentialsIfUnchanged(
  1152. s.ctx,
  1153. account.ID,
  1154. observed.Credentials,
  1155. observed.ProxyID,
  1156. map[string]any{
  1157. "access_token": "provider-access",
  1158. "refresh_token": "provider-refresh",
  1159. "_token_version": int64(22),
  1160. },
  1161. )
  1162. s.Require().NoError(err)
  1163. s.Require().False(applied)
  1164. got, err := s.repo.GetByID(s.ctx, account.ID)
  1165. s.Require().NoError(err)
  1166. s.Require().Equal("reauthorized-refresh", got.GetGrokRefreshToken())
  1167. s.Require().Empty(cacheRecorder.setAccounts)
  1168. var outboxCount int
  1169. err = scanSingleRow(
  1170. s.ctx,
  1171. s.repo.sql,
  1172. "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
  1173. []any{service.SchedulerOutboxEventAccountChanged, account.ID},
  1174. &outboxCount,
  1175. )
  1176. s.Require().NoError(err)
  1177. s.Require().Zero(outboxCount)
  1178. }
  1179. func (s *AccountRepoSuite) TestGrokOAuthConditionalMutation_DetachesBoundedSnapshotSync() {
  1180. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1181. Name: "grok-conditional-detached-sync",
  1182. Platform: service.PlatformGrok,
  1183. Type: service.AccountTypeOAuth,
  1184. Status: service.StatusActive,
  1185. Schedulable: true,
  1186. Credentials: map[string]any{"access_token": "observed"},
  1187. })
  1188. observed, err := s.repo.GetByID(s.ctx, account.ID)
  1189. s.Require().NoError(err)
  1190. ctx, cancel := context.WithCancel(context.Background())
  1191. cacheRecorder := &schedulerCacheRecorder{}
  1192. repo := newAccountRepositoryWithSQL(s.client, &cancelAfterAtomicMutationSQLExecutor{
  1193. sqlExecutor: s.repo.sql,
  1194. cancel: cancel,
  1195. }, cacheRecorder)
  1196. applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
  1197. ctx,
  1198. account.ID,
  1199. observed.Credentials,
  1200. "missing refresh token",
  1201. )
  1202. s.Require().NoError(err)
  1203. s.Require().True(applied)
  1204. s.Require().ErrorIs(ctx.Err(), context.Canceled)
  1205. s.Require().Len(cacheRecorder.setAccounts, 1)
  1206. s.Require().NoError(cacheRecorder.setCtxErr, "immediate scheduler propagation must use a bounded detached context")
  1207. }
  1208. func TestGrokOAuthConditionalMutationRollsBackWhenOutboxInsertFails(t *testing.T) {
  1209. client := testEntClient(t)
  1210. account := mustCreateAccount(t, client, &service.Account{
  1211. Name: "grok-conditional-atomic-outbox-failure",
  1212. Platform: service.PlatformGrok,
  1213. Type: service.AccountTypeOAuth,
  1214. Status: service.StatusActive,
  1215. Schedulable: true,
  1216. Credentials: map[string]any{"access_token": "observed"},
  1217. })
  1218. t.Cleanup(func() {
  1219. _, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM scheduler_outbox WHERE account_id = $1", account.ID)
  1220. _ = client.Account.DeleteOneID(account.ID).Exec(context.Background())
  1221. })
  1222. repo := newAccountRepositoryWithSQL(client, &failAtomicSchedulerOutboxSQLExecutor{sqlExecutor: integrationDB}, nil)
  1223. applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
  1224. context.Background(),
  1225. account.ID,
  1226. account.Credentials,
  1227. "missing refresh token",
  1228. )
  1229. require.Error(t, err)
  1230. require.False(t, applied)
  1231. got, readErr := repo.GetByID(context.Background(), account.ID)
  1232. require.NoError(t, readErr)
  1233. require.Equal(t, service.StatusActive, got.Status)
  1234. require.True(t, got.Schedulable)
  1235. require.Empty(t, got.ErrorMessage)
  1236. var outboxCount int
  1237. require.NoError(t, integrationDB.QueryRowContext(
  1238. context.Background(),
  1239. "SELECT COUNT(*) FROM scheduler_outbox WHERE account_id = $1",
  1240. account.ID,
  1241. ).Scan(&outboxCount))
  1242. require.Zero(t, outboxCount)
  1243. }
  1244. func (s *AccountRepoSuite) TestUpdateErrorStatusUnschedulesAccount() {
  1245. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-update-err", Status: service.StatusActive, Schedulable: true})
  1246. account.Status = service.StatusError
  1247. account.ErrorMessage = "token revoked"
  1248. account.Schedulable = true
  1249. s.Require().NoError(s.repo.Update(s.ctx, account))
  1250. got, err := s.repo.GetByID(s.ctx, account.ID)
  1251. s.Require().NoError(err)
  1252. s.Require().Equal(service.StatusError, got.Status)
  1253. s.Require().Equal("token revoked", got.ErrorMessage)
  1254. s.Require().False(got.Schedulable)
  1255. }
  1256. func (s *AccountRepoSuite) TestClearError_SyncSchedulerSnapshotOnRecovery() {
  1257. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1258. Name: "acc-clear-err",
  1259. Status: service.StatusError,
  1260. ErrorMessage: "temporary error",
  1261. })
  1262. cacheRecorder := &schedulerCacheRecorder{}
  1263. s.repo.schedulerCache = cacheRecorder
  1264. s.Require().NoError(s.repo.ClearError(s.ctx, account.ID))
  1265. got, err := s.repo.GetByID(s.ctx, account.ID)
  1266. s.Require().NoError(err)
  1267. s.Require().Equal(service.StatusActive, got.Status)
  1268. s.Require().Empty(got.ErrorMessage)
  1269. s.Require().Len(cacheRecorder.setAccounts, 1)
  1270. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  1271. s.Require().Equal(service.StatusActive, cacheRecorder.setAccounts[0].Status)
  1272. }
  1273. // --- UpdateSessionWindow ---
  1274. func (s *AccountRepoSuite) TestUpdateSessionWindow() {
  1275. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-win"})
  1276. start := time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC)
  1277. end := time.Date(2025, 6, 15, 15, 0, 0, 0, time.UTC)
  1278. s.Require().NoError(s.repo.UpdateSessionWindow(s.ctx, account.ID, &start, &end, "active"))
  1279. got, err := s.repo.GetByID(s.ctx, account.ID)
  1280. s.Require().NoError(err)
  1281. s.Require().NotNil(got.SessionWindowStart)
  1282. s.Require().NotNil(got.SessionWindowEnd)
  1283. s.Require().Equal("active", got.SessionWindowStatus)
  1284. }
  1285. // --- UpdateExtra ---
  1286. func (s *AccountRepoSuite) TestUpdateExtra_MergesFields() {
  1287. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1288. Name: "acc-extra",
  1289. Extra: map[string]any{"a": "1"},
  1290. })
  1291. s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{"b": "2"}), "UpdateExtra")
  1292. got, err := s.repo.GetByID(s.ctx, account.ID)
  1293. s.Require().NoError(err, "GetByID")
  1294. s.Require().Equal("1", got.Extra["a"])
  1295. s.Require().Equal("2", got.Extra["b"])
  1296. }
  1297. func (s *AccountRepoSuite) TestUpdateExtra_EmptyUpdates() {
  1298. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-extra-empty"})
  1299. s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{}))
  1300. }
  1301. func (s *AccountRepoSuite) TestUpdateExtra_NilExtra() {
  1302. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-nil-extra", Extra: nil})
  1303. s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{"key": "val"}))
  1304. got, err := s.repo.GetByID(s.ctx, account.ID)
  1305. s.Require().NoError(err)
  1306. s.Require().Equal("val", got.Extra["key"])
  1307. }
  1308. func (s *AccountRepoSuite) TestUpdateExtra_SchedulerNeutralSkipsOutboxAndSyncsFreshSnapshot() {
  1309. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1310. Name: "acc-extra-neutral",
  1311. Platform: service.PlatformOpenAI,
  1312. Extra: map[string]any{"codex_usage_updated_at": "old"},
  1313. })
  1314. cacheRecorder := &schedulerCacheRecorder{
  1315. accounts: map[int64]*service.Account{
  1316. account.ID: {
  1317. ID: account.ID,
  1318. Platform: account.Platform,
  1319. Status: service.StatusDisabled,
  1320. Extra: map[string]any{
  1321. "codex_usage_updated_at": "old",
  1322. },
  1323. },
  1324. },
  1325. }
  1326. s.repo.schedulerCache = cacheRecorder
  1327. updates := map[string]any{
  1328. "codex_usage_updated_at": "2026-03-11T10:00:00Z",
  1329. "codex_5h_used_percent": 88.5,
  1330. "session_window_utilization": 0.42,
  1331. }
  1332. s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, updates))
  1333. got, err := s.repo.GetByID(s.ctx, account.ID)
  1334. s.Require().NoError(err)
  1335. s.Require().Equal("2026-03-11T10:00:00Z", got.Extra["codex_usage_updated_at"])
  1336. s.Require().Equal(88.5, got.Extra["codex_5h_used_percent"])
  1337. s.Require().Equal(0.42, got.Extra["session_window_utilization"])
  1338. var outboxCount int
  1339. s.Require().NoError(scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &outboxCount))
  1340. s.Require().Zero(outboxCount)
  1341. s.Require().Len(cacheRecorder.setAccounts, 1)
  1342. s.Require().NotNil(cacheRecorder.accounts[account.ID])
  1343. s.Require().Equal(service.StatusActive, cacheRecorder.accounts[account.ID].Status)
  1344. s.Require().Equal("2026-03-11T10:00:00Z", cacheRecorder.accounts[account.ID].Extra["codex_usage_updated_at"])
  1345. }
  1346. func (s *AccountRepoSuite) TestUpdateExtra_ExhaustedCodexSnapshotSyncsSchedulerCache() {
  1347. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1348. Name: "acc-extra-codex-exhausted",
  1349. Platform: service.PlatformOpenAI,
  1350. Type: service.AccountTypeOAuth,
  1351. Extra: map[string]any{},
  1352. })
  1353. cacheRecorder := &schedulerCacheRecorder{}
  1354. s.repo.schedulerCache = cacheRecorder
  1355. _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  1356. s.Require().NoError(err)
  1357. s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{
  1358. "codex_7d_used_percent": 100.0,
  1359. "codex_7d_reset_at": "2026-03-12T13:00:00Z",
  1360. "codex_7d_reset_after_seconds": 86400,
  1361. }))
  1362. var count int
  1363. err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
  1364. s.Require().NoError(err)
  1365. s.Require().Equal(0, count)
  1366. s.Require().Len(cacheRecorder.setAccounts, 1)
  1367. s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
  1368. s.Require().Equal(service.StatusActive, cacheRecorder.setAccounts[0].Status)
  1369. s.Require().Equal(100.0, cacheRecorder.setAccounts[0].Extra["codex_7d_used_percent"])
  1370. }
  1371. func (s *AccountRepoSuite) TestUpdateExtra_SchedulerRelevantStillEnqueuesOutbox() {
  1372. account := mustCreateAccount(s.T(), s.client, &service.Account{
  1373. Name: "acc-extra-mixed",
  1374. Platform: service.PlatformAntigravity,
  1375. Extra: map[string]any{},
  1376. })
  1377. _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
  1378. s.Require().NoError(err)
  1379. s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{
  1380. "mixed_scheduling": true,
  1381. "codex_usage_updated_at": "2026-03-11T10:00:00Z",
  1382. }))
  1383. var count int
  1384. err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
  1385. s.Require().NoError(err)
  1386. s.Require().Equal(1, count)
  1387. }
  1388. // --- GetByCRSAccountID ---
  1389. func (s *AccountRepoSuite) TestGetByCRSAccountID() {
  1390. crsID := "crs-12345"
  1391. mustCreateAccount(s.T(), s.client, &service.Account{
  1392. Name: "acc-crs",
  1393. Extra: map[string]any{"crs_account_id": crsID},
  1394. })
  1395. got, err := s.repo.GetByCRSAccountID(s.ctx, crsID)
  1396. s.Require().NoError(err)
  1397. s.Require().NotNil(got)
  1398. s.Require().Equal("acc-crs", got.Name)
  1399. }
  1400. func (s *AccountRepoSuite) TestGetByCRSAccountID_NotFound() {
  1401. got, err := s.repo.GetByCRSAccountID(s.ctx, "non-existent")
  1402. s.Require().NoError(err)
  1403. s.Require().Nil(got)
  1404. }
  1405. func (s *AccountRepoSuite) TestGetByCRSAccountID_EmptyString() {
  1406. got, err := s.repo.GetByCRSAccountID(s.ctx, "")
  1407. s.Require().NoError(err)
  1408. s.Require().Nil(got)
  1409. }
  1410. // TestGetByCRSAccountID_ExcludesSparkShadow 验证外审第7轮 P1:即便 spark 影子的 Extra 被误写入
  1411. // crs_account_id,CRS 查询也绝不能命中影子(否则会被当普通账号更新而覆盖 type/credentials/proxy)。
  1412. func (s *AccountRepoSuite) TestGetByCRSAccountID_ExcludesSparkShadow() {
  1413. crsID := "crs-shadow-only-99"
  1414. parent := mustCreateAccount(s.T(), s.client, &service.Account{
  1415. Name: "crs-mother", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
  1416. })
  1417. mustCreateAccount(s.T(), s.client, &service.Account{
  1418. Name: "crs-shadow", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
  1419. ParentAccountID: &parent.ID,
  1420. QuotaDimension: service.QuotaDimensionSpark,
  1421. Extra: map[string]any{"crs_account_id": crsID},
  1422. })
  1423. got, err := s.repo.GetByCRSAccountID(s.ctx, crsID)
  1424. s.Require().NoError(err)
  1425. s.Require().Nil(got, "spark 影子即便带 crs_account_id 也不应被 CRS 命中")
  1426. }
  1427. // TestListCRSAccountIDs_ExcludesSparkShadow 验证外审第7轮 P1:影子的 crs_account_id 不应进入
  1428. // CRS 同步映射(否则后续 CRS 同步会把影子当普通账号更新)。
  1429. func (s *AccountRepoSuite) TestListCRSAccountIDs_ExcludesSparkShadow() {
  1430. parent := mustCreateAccount(s.T(), s.client, &service.Account{
  1431. Name: "crs-list-mother", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
  1432. })
  1433. shadowCRSID := "crs-list-shadow-77"
  1434. mustCreateAccount(s.T(), s.client, &service.Account{
  1435. Name: "crs-list-shadow", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
  1436. ParentAccountID: &parent.ID,
  1437. QuotaDimension: service.QuotaDimensionSpark,
  1438. Extra: map[string]any{"crs_account_id": shadowCRSID},
  1439. })
  1440. ids, err := s.repo.ListCRSAccountIDs(s.ctx)
  1441. s.Require().NoError(err)
  1442. _, ok := ids[shadowCRSID]
  1443. s.Require().False(ok, "影子的 crs_account_id 不应进入 CRS 映射")
  1444. }
  1445. // --- BulkUpdate ---
  1446. func (s *AccountRepoSuite) TestBulkUpdate() {
  1447. a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk1", Priority: 1})
  1448. a2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk2", Priority: 1})
  1449. newPriority := 99
  1450. affected, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID, a2.ID}, service.AccountBulkUpdate{
  1451. Priority: &newPriority,
  1452. })
  1453. s.Require().NoError(err)
  1454. s.Require().GreaterOrEqual(affected, int64(1), "expected at least one affected row")
  1455. got1, _ := s.repo.GetByID(s.ctx, a1.ID)
  1456. got2, _ := s.repo.GetByID(s.ctx, a2.ID)
  1457. s.Require().Equal(99, got1.Priority)
  1458. s.Require().Equal(99, got2.Priority)
  1459. }
  1460. func (s *AccountRepoSuite) TestBulkUpdate_MergeCredentials() {
  1461. a1 := mustCreateAccount(s.T(), s.client, &service.Account{
  1462. Name: "bulk-cred",
  1463. Credentials: map[string]any{"existing": "value"},
  1464. })
  1465. _, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{
  1466. Credentials: map[string]any{"new_key": "new_value"},
  1467. })
  1468. s.Require().NoError(err)
  1469. got, _ := s.repo.GetByID(s.ctx, a1.ID)
  1470. s.Require().Equal("value", got.Credentials["existing"])
  1471. s.Require().Equal("new_value", got.Credentials["new_key"])
  1472. }
  1473. func (s *AccountRepoSuite) TestBulkUpdate_MergeExtra() {
  1474. a1 := mustCreateAccount(s.T(), s.client, &service.Account{
  1475. Name: "bulk-extra",
  1476. Extra: map[string]any{"existing": "val"},
  1477. })
  1478. _, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{
  1479. Extra: map[string]any{"new_key": "new_val"},
  1480. })
  1481. s.Require().NoError(err)
  1482. got, _ := s.repo.GetByID(s.ctx, a1.ID)
  1483. s.Require().Equal("val", got.Extra["existing"])
  1484. s.Require().Equal("new_val", got.Extra["new_key"])
  1485. }
  1486. func (s *AccountRepoSuite) TestBulkUpdate_EmptyIDs() {
  1487. affected, err := s.repo.BulkUpdate(s.ctx, []int64{}, service.AccountBulkUpdate{})
  1488. s.Require().NoError(err)
  1489. s.Require().Zero(affected)
  1490. }
  1491. func (s *AccountRepoSuite) TestBulkUpdate_EmptyUpdates() {
  1492. a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-empty"})
  1493. affected, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{})
  1494. s.Require().NoError(err)
  1495. s.Require().Zero(affected)
  1496. }
  1497. func idsOfAccounts(accounts []service.Account) []int64 {
  1498. out := make([]int64, 0, len(accounts))
  1499. for i := range accounts {
  1500. out = append(out, accounts[i].ID)
  1501. }
  1502. return out
  1503. }