| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728 |
- //go:build integration
- package repository
- import (
- "context"
- "database/sql"
- "strings"
- "testing"
- "time"
- dbent "github.com/Wei-Shaw/sub2api/ent"
- "github.com/Wei-Shaw/sub2api/ent/accountgroup"
- "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/stretchr/testify/require"
- "github.com/stretchr/testify/suite"
- )
- type AccountRepoSuite struct {
- suite.Suite
- ctx context.Context
- client *dbent.Client
- repo *accountRepository
- }
- type schedulerCacheRecorder struct {
- setAccounts []*service.Account
- deleteIDs []int64
- accounts map[int64]*service.Account
- setCtxErr error
- }
- func (s *schedulerCacheRecorder) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
- return nil, false, nil
- }
- func (s *schedulerCacheRecorder) CaptureBucketWriteToken(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
- return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
- }
- func (s *schedulerCacheRecorder) SetSnapshot(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) error {
- return nil
- }
- func (s *schedulerCacheRecorder) RetireBucket(ctx context.Context, bucket service.SchedulerBucket) error {
- return nil
- }
- func (s *schedulerCacheRecorder) ReopenBucket(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
- return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
- }
- func (s *schedulerCacheRecorder) TryAcquireGroupLifecycleLease(_ context.Context, groupID int64, _ time.Duration) (service.SchedulerGroupLifecycleLease, bool, error) {
- return service.SchedulerGroupLifecycleLease{GroupID: groupID, OwnerToken: "scheduler-cache-recorder"}, true, nil
- }
- func (s *schedulerCacheRecorder) ReleaseGroupLifecycleLease(context.Context, service.SchedulerGroupLifecycleLease) error {
- return nil
- }
- func (s *schedulerCacheRecorder) GetAccount(ctx context.Context, accountID int64) (*service.Account, error) {
- if s.accounts == nil {
- return nil, nil
- }
- return s.accounts[accountID], nil
- }
- func (s *schedulerCacheRecorder) SetAccount(ctx context.Context, account *service.Account) error {
- s.setCtxErr = ctx.Err()
- s.setAccounts = append(s.setAccounts, account)
- if s.accounts == nil {
- s.accounts = make(map[int64]*service.Account)
- }
- if account != nil {
- s.accounts[account.ID] = account
- }
- return nil
- }
- type failAtomicSchedulerOutboxSQLExecutor struct {
- sqlExecutor
- }
- func (e *failAtomicSchedulerOutboxSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
- if strings.Contains(query, "WITH updated AS") && strings.Contains(query, "INSERT INTO scheduler_outbox") && len(args) > 0 {
- args = append([]any(nil), args...)
- args[len(args)-1] = nil // event_type is NOT NULL; the whole statement must roll back.
- }
- return e.sqlExecutor.ExecContext(ctx, query, args...)
- }
- type cancelAfterAtomicMutationSQLExecutor struct {
- sqlExecutor
- cancel context.CancelFunc
- }
- func (e *cancelAfterAtomicMutationSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
- result, err := e.sqlExecutor.ExecContext(ctx, query, args...)
- if err == nil && strings.Contains(query, "WITH updated AS") && strings.Contains(query, "INSERT INTO scheduler_outbox") {
- e.cancel()
- }
- return result, err
- }
- func (s *schedulerCacheRecorder) DeleteAccount(ctx context.Context, accountID int64) error {
- s.deleteIDs = append(s.deleteIDs, accountID)
- if s.accounts != nil {
- delete(s.accounts, accountID)
- }
- return nil
- }
- func (s *schedulerCacheRecorder) UpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
- return nil
- }
- func (s *schedulerCacheRecorder) TryLockBucket(ctx context.Context, bucket service.SchedulerBucket, ttl time.Duration) (bool, error) {
- return true, nil
- }
- func (s *schedulerCacheRecorder) UnlockBucket(ctx context.Context, bucket service.SchedulerBucket) error {
- return nil
- }
- func (s *schedulerCacheRecorder) ListBuckets(ctx context.Context) ([]service.SchedulerBucket, error) {
- return nil, nil
- }
- func (s *schedulerCacheRecorder) GetOutboxWatermark(ctx context.Context) (int64, error) {
- return 0, nil
- }
- func (s *schedulerCacheRecorder) SetOutboxWatermark(ctx context.Context, id int64) error {
- return nil
- }
- func (s *AccountRepoSuite) SetupTest() {
- s.ctx = context.Background()
- tx := testEntTx(s.T())
- s.client = tx.Client()
- s.repo = newAccountRepositoryWithSQL(s.client, tx, nil)
- }
- func TestAccountRepoSuite(t *testing.T) {
- suite.Run(t, new(AccountRepoSuite))
- }
- // --- Create / GetByID / Update / Delete ---
- func (s *AccountRepoSuite) TestCreate() {
- account := &service.Account{
- Name: "test-create",
- Platform: service.PlatformAnthropic,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{},
- Extra: map[string]any{},
- Concurrency: 3,
- Priority: 50,
- Schedulable: true,
- }
- err := s.repo.Create(s.ctx, account)
- s.Require().NoError(err, "Create")
- s.Require().NotZero(account.ID, "expected ID to be set")
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err, "GetByID")
- s.Require().Equal("test-create", got.Name)
- }
- func (s *AccountRepoSuite) TestGetByID_NotFound() {
- _, err := s.repo.GetByID(s.ctx, 999999)
- s.Require().Error(err, "expected error for non-existent ID")
- }
- func (s *AccountRepoSuite) TestUpdate() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "original"})
- account.Name = "updated"
- err := s.repo.Update(s.ctx, account)
- s.Require().NoError(err, "Update")
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err, "GetByID after update")
- s.Require().Equal("updated", got.Name)
- }
- func (s *AccountRepoSuite) TestUpdate_SyncSchedulerSnapshotOnDisabled() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "sync-update", Status: service.StatusActive, Schedulable: true})
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- account.Status = service.StatusDisabled
- err := s.repo.Update(s.ctx, account)
- s.Require().NoError(err, "Update")
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().Equal(service.StatusDisabled, cacheRecorder.setAccounts[0].Status)
- }
- func (s *AccountRepoSuite) TestUpdate_SyncSchedulerSnapshotOnCredentialsChange() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "sync-credentials-update",
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "gpt-5": "gpt-5.1",
- },
- },
- })
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- account.Credentials = map[string]any{
- "model_mapping": map[string]any{
- "gpt-5": "gpt-5.2",
- },
- }
- err := s.repo.Update(s.ctx, account)
- s.Require().NoError(err, "Update")
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- mapping, ok := cacheRecorder.setAccounts[0].Credentials["model_mapping"].(map[string]any)
- s.Require().True(ok)
- s.Require().Equal("gpt-5.2", mapping["gpt-5"])
- }
- func (s *AccountRepoSuite) TestUpdateCredentials_SyncsSnapshotAndDurableOutbox() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "sync-refresh-credentials",
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{"access_token": "old-token"},
- })
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{"access_token": "new-token"}))
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal("new-token", cacheRecorder.setAccounts[0].GetCredential("access_token"))
- var outboxCount int
- err = scanSingleRow(
- s.ctx,
- s.repo.sql,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- []any{service.SchedulerOutboxEventAccountChanged, account.ID},
- &outboxCount,
- )
- s.Require().NoError(err)
- s.Require().Equal(1, outboxCount)
- }
- func (s *AccountRepoSuite) TestDelete() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "to-delete"})
- err := s.repo.Delete(s.ctx, account.ID)
- s.Require().NoError(err, "Delete")
- _, err = s.repo.GetByID(s.ctx, account.ID)
- s.Require().Error(err, "expected error after delete")
- }
- func (s *AccountRepoSuite) TestDelete_RemovesSchedulerAccountSnapshot() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "to-delete-cache"})
- cacheRecorder := &schedulerCacheRecorder{
- accounts: map[int64]*service.Account{
- account.ID: {
- ID: account.ID,
- Name: account.Name,
- Status: service.StatusActive,
- Schedulable: true,
- },
- },
- }
- s.repo.schedulerCache = cacheRecorder
- err := s.repo.Delete(s.ctx, account.ID)
- s.Require().NoError(err, "Delete")
- s.Require().Equal([]int64{account.ID}, cacheRecorder.deleteIDs)
- s.Require().NotContains(cacheRecorder.accounts, account.ID)
- }
- func (s *AccountRepoSuite) TestDelete_WithGroupBindings() {
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-del"})
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-del"})
- mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
- err := s.repo.Delete(s.ctx, account.ID)
- s.Require().NoError(err, "Delete should cascade remove bindings")
- count, err := s.client.AccountGroup.Query().Where(accountgroup.AccountIDEQ(account.ID)).Count(s.ctx)
- s.Require().NoError(err)
- s.Require().Zero(count, "expected bindings to be removed")
- }
- // --- List / ListWithFilters ---
- func (s *AccountRepoSuite) TestList() {
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc1"})
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc2"})
- accounts, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10})
- s.Require().NoError(err, "List")
- s.Require().Len(accounts, 2)
- s.Require().Equal(int64(2), page.Total)
- }
- func (s *AccountRepoSuite) TestListOAuthRefreshCandidatePage_GrokCursorAndExclusions() {
- now := time.Now().UTC()
- valid1 := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-oauth-page-1",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{
- "access_token": "access-1",
- "refresh_token": "refresh-1",
- "expires_at": now.Add(30 * time.Minute).Format(time.RFC3339),
- },
- })
- unschedulable := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-oauth-unschedulable-excluded",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{"refresh_token": "refresh-unschedulable"},
- })
- s.Require().NoError(s.client.Account.UpdateOneID(unschedulable.ID).SetSchedulable(false).Exec(s.ctx))
- mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-api-key-excluded",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeAPIKey,
- Status: service.StatusActive,
- Credentials: map[string]any{
- "api_key": "api-key",
- "refresh_token": "must-not-make-api-key-eligible",
- },
- })
- valid2 := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-oauth-page-2",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{"refresh_token": "refresh-2"},
- })
- mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-oauth-blank-refresh-excluded",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{"refresh_token": " "},
- })
- valid3 := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-oauth-page-3",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{"refresh_token": "refresh-3"},
- })
- cooldown := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-oauth-retry-cooldown-excluded",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{"refresh_token": "refresh-cooldown"},
- })
- s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, cooldown.ID, now.Add(10*time.Minute), "token refresh retry exhausted: timeout"))
- mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "openai-oauth-excluded",
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Credentials: map[string]any{"refresh_token": "refresh-openai"},
- })
- options := service.OAuthRefreshPageOptions{
- Platforms: []string{service.PlatformGrok},
- Limit: 2,
- ActiveOnly: true,
- RequireRefreshToken: true,
- ExcludeRetryCooldown: true,
- }
- firstPage, err := s.repo.ListOAuthRefreshCandidatePage(s.ctx, options)
- s.Require().NoError(err)
- first := firstPage.Accounts
- s.Require().Len(first, 2)
- s.Require().Equal([]int64{valid1.ID, valid2.ID}, []int64{first[0].ID, first[1].ID})
- s.Require().NotContains([]int64{first[0].ID, first[1].ID}, unschedulable.ID)
- options.AfterID = first[len(first)-1].ID
- secondPage, err := s.repo.ListOAuthRefreshCandidatePage(s.ctx, options)
- s.Require().NoError(err)
- second := secondPage.Accounts
- s.Require().Len(second, 1)
- s.Require().Equal(valid3.ID, second[0].ID)
- s.Require().NotContains([]int64{first[0].ID, first[1].ID}, second[0].ID)
- }
- func (s *AccountRepoSuite) TestListWithFilters() {
- tests := []struct {
- name string
- setup func(client *dbent.Client)
- platform string
- accType string
- status string
- search string
- groupID int64
- privacyMode string
- wantCount int
- validate func(accounts []service.Account)
- }{
- {
- name: "filter_by_platform",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic})
- mustCreateAccount(s.T(), client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI})
- },
- platform: service.PlatformOpenAI,
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal(service.PlatformOpenAI, accounts[0].Platform)
- },
- },
- {
- name: "filter_by_type",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "t1", Type: service.AccountTypeOAuth})
- mustCreateAccount(s.T(), client, &service.Account{Name: "t2", Type: service.AccountTypeAPIKey})
- },
- accType: service.AccountTypeAPIKey,
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal(service.AccountTypeAPIKey, accounts[0].Type)
- },
- },
- {
- name: "filter_by_status",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "s1", Status: service.StatusActive})
- mustCreateAccount(s.T(), client, &service.Account{Name: "s2", Status: service.StatusDisabled})
- },
- status: service.StatusDisabled,
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal(service.StatusDisabled, accounts[0].Status)
- },
- },
- {
- name: "filter_by_status_active_excludes_runtime_blocked_accounts",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "active-normal", Status: service.StatusActive})
- rateLimited := mustCreateAccount(s.T(), client, &service.Account{Name: "active-rate-limited", Status: service.StatusActive})
- err := client.Account.UpdateOneID(rateLimited.ID).
- SetRateLimitResetAt(time.Now().Add(10 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive})
- err = client.Account.UpdateOneID(tempUnsched.ID).
- SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- unsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-unsched", Status: service.StatusActive})
- err = client.Account.UpdateOneID(unsched.ID).
- SetSchedulable(false).
- Exec(context.Background())
- s.Require().NoError(err)
- },
- status: service.StatusActive,
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal("active-normal", accounts[0].Name)
- },
- },
- {
- name: "filter_by_status_unschedulable_excludes_rate_limited_and_temp_unschedulable",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "active-normal", Status: service.StatusActive, Schedulable: true})
- unsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-unsched", Status: service.StatusActive})
- err := client.Account.UpdateOneID(unsched.ID).
- SetSchedulable(false).
- Exec(context.Background())
- s.Require().NoError(err)
- rateLimited := mustCreateAccount(s.T(), client, &service.Account{Name: "active-rate-limited", Status: service.StatusActive})
- err = client.Account.UpdateOneID(rateLimited.ID).
- SetSchedulable(false).
- SetRateLimitResetAt(time.Now().Add(10 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive})
- err = client.Account.UpdateOneID(tempUnsched.ID).
- SetSchedulable(false).
- SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- },
- status: "unschedulable",
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal("active-unsched", accounts[0].Name)
- },
- },
- {
- name: "filter_by_status_rate_limited_excludes_temp_unschedulable",
- setup: func(client *dbent.Client) {
- rateLimited := mustCreateAccount(s.T(), client, &service.Account{Name: "active-rate-limited", Status: service.StatusActive})
- err := client.Account.UpdateOneID(rateLimited.ID).
- SetRateLimitResetAt(time.Now().Add(10 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive})
- err = client.Account.UpdateOneID(tempUnsched.ID).
- SetRateLimitResetAt(time.Now().Add(20 * time.Minute)).
- SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- },
- status: "rate_limited",
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal("active-rate-limited", accounts[0].Name)
- },
- },
- {
- name: "filter_by_status_temp_unschedulable_excludes_manually_unschedulable",
- setup: func(client *dbent.Client) {
- tempUnsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-temp-unsched", Status: service.StatusActive, Schedulable: true})
- err := client.Account.UpdateOneID(tempUnsched.ID).
- SetTempUnschedulableUntil(time.Now().Add(15 * time.Minute)).
- Exec(context.Background())
- s.Require().NoError(err)
- unsched := mustCreateAccount(s.T(), client, &service.Account{Name: "active-unsched", Status: service.StatusActive})
- err = client.Account.UpdateOneID(unsched.ID).
- SetSchedulable(false).
- Exec(context.Background())
- s.Require().NoError(err)
- },
- status: "temp_unschedulable",
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal("active-temp-unsched", accounts[0].Name)
- },
- },
- {
- name: "filter_by_search",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "alpha-account"})
- mustCreateAccount(s.T(), client, &service.Account{Name: "beta-account"})
- },
- search: "alpha",
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Contains(accounts[0].Name, "alpha")
- },
- },
- {
- name: "filter_by_ungrouped",
- setup: func(client *dbent.Client) {
- group := mustCreateGroup(s.T(), client, &service.Group{Name: "g-ungrouped"})
- grouped := mustCreateAccount(s.T(), client, &service.Account{Name: "grouped-account"})
- mustCreateAccount(s.T(), client, &service.Account{Name: "ungrouped-account"})
- mustBindAccountToGroup(s.T(), client, grouped.ID, group.ID, 1)
- },
- groupID: service.AccountListGroupUngrouped,
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal("ungrouped-account", accounts[0].Name)
- s.Require().Empty(accounts[0].GroupIDs)
- },
- },
- {
- name: "filter_by_privacy_mode",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-ok", Extra: map[string]any{"privacy_mode": service.PrivacyModeTrainingOff}})
- mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-fail", Extra: map[string]any{"privacy_mode": service.PrivacyModeFailed}})
- },
- privacyMode: service.PrivacyModeTrainingOff,
- wantCount: 1,
- validate: func(accounts []service.Account) {
- s.Require().Equal("privacy-ok", accounts[0].Name)
- },
- },
- {
- name: "filter_by_privacy_mode_unset",
- setup: func(client *dbent.Client) {
- mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-unset", Extra: nil})
- mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-empty", Extra: map[string]any{"privacy_mode": ""}})
- mustCreateAccount(s.T(), client, &service.Account{Name: "privacy-set", Extra: map[string]any{"privacy_mode": service.PrivacyModeTrainingOff}})
- },
- privacyMode: service.AccountPrivacyModeUnsetFilter,
- wantCount: 2,
- validate: func(accounts []service.Account) {
- names := []string{accounts[0].Name, accounts[1].Name}
- s.ElementsMatch([]string{"privacy-unset", "privacy-empty"}, names)
- },
- },
- }
- for _, tt := range tests {
- s.Run(tt.name, func() {
- // 每个 case 重新获取隔离资源
- tx := testEntTx(s.T())
- client := tx.Client()
- repo := newAccountRepositoryWithSQL(client, tx, nil)
- ctx := context.Background()
- tt.setup(client)
- accounts, page, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, tt.platform, tt.accType, tt.status, tt.search, tt.groupID, tt.privacyMode)
- s.Require().NoError(err)
- s.Require().Len(accounts, tt.wantCount)
- // Regression guard for issue #3601: when the whole result set fits on a single page,
- // pagination.Total must match len(items). A mismatch means the Count query was applied
- // against different predicates than the list query — the exact symptom reported.
- s.Require().NotNil(page)
- s.Require().Equal(int64(tt.wantCount), page.Total, "total must match items on single page")
- if tt.validate != nil {
- tt.validate(accounts)
- }
- })
- }
- }
- // --- ListByGroup / ListActive / ListByPlatform ---
- func (s *AccountRepoSuite) TestListByGroup() {
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-list"})
- acc1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Status: service.StatusActive})
- acc2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Status: service.StatusActive})
- mustBindAccountToGroup(s.T(), s.client, acc1.ID, group.ID, 2)
- mustBindAccountToGroup(s.T(), s.client, acc2.ID, group.ID, 1)
- accounts, err := s.repo.ListByGroup(s.ctx, group.ID)
- s.Require().NoError(err, "ListByGroup")
- s.Require().Len(accounts, 2)
- // Should be ordered by priority
- s.Require().Equal(acc2.ID, accounts[0].ID, "expected acc2 first (priority=1)")
- }
- func (s *AccountRepoSuite) TestListActive() {
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "active1", Status: service.StatusActive})
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "inactive1", Status: service.StatusDisabled})
- accounts, err := s.repo.ListActive(s.ctx)
- s.Require().NoError(err, "ListActive")
- s.Require().Len(accounts, 1)
- s.Require().Equal("active1", accounts[0].Name)
- }
- func (s *AccountRepoSuite) TestListByPlatform() {
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "p1", Platform: service.PlatformAnthropic, Status: service.StatusActive})
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "p2", Platform: service.PlatformOpenAI, Status: service.StatusActive})
- accounts, err := s.repo.ListByPlatform(s.ctx, service.PlatformAnthropic)
- s.Require().NoError(err, "ListByPlatform")
- s.Require().Len(accounts, 1)
- s.Require().Equal(service.PlatformAnthropic, accounts[0].Platform)
- }
- // --- Preload and VirtualFields ---
- func (s *AccountRepoSuite) TestPreload_And_VirtualFields() {
- proxy := mustCreateProxy(s.T(), s.client, &service.Proxy{Name: "p1"})
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g1"})
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc1",
- ProxyID: &proxy.ID,
- })
- mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err, "GetByID")
- s.Require().NotNil(got.Proxy, "expected Proxy preload")
- s.Require().Equal(proxy.ID, got.Proxy.ID)
- s.Require().Len(got.GroupIDs, 1, "expected GroupIDs to be populated")
- s.Require().Equal(group.ID, got.GroupIDs[0])
- s.Require().Len(got.Groups, 1, "expected Groups to be populated")
- s.Require().Equal(group.ID, got.Groups[0].ID)
- accounts, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, "", "", "", "acc", 0, "")
- s.Require().NoError(err, "ListWithFilters")
- s.Require().Equal(int64(1), page.Total)
- s.Require().Len(accounts, 1)
- s.Require().NotNil(accounts[0].Proxy, "expected Proxy preload in list")
- s.Require().Equal(proxy.ID, accounts[0].Proxy.ID)
- s.Require().Len(accounts[0].GroupIDs, 1, "expected GroupIDs in list")
- s.Require().Equal(group.ID, accounts[0].GroupIDs[0])
- }
- // --- GroupBinding / AddToGroup / RemoveFromGroup / BindGroups / GetGroups ---
- func (s *AccountRepoSuite) TestGroupBinding_And_BindGroups() {
- g1 := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g1"})
- g2 := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g2"})
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc"})
- s.Require().NoError(s.repo.AddToGroup(s.ctx, account.ID, g1.ID, 10), "AddToGroup")
- groups, err := s.repo.GetGroups(s.ctx, account.ID)
- s.Require().NoError(err, "GetGroups")
- s.Require().Len(groups, 1, "expected 1 group")
- s.Require().Equal(g1.ID, groups[0].ID)
- s.Require().NoError(s.repo.RemoveFromGroup(s.ctx, account.ID, g1.ID), "RemoveFromGroup")
- groups, err = s.repo.GetGroups(s.ctx, account.ID)
- s.Require().NoError(err, "GetGroups after remove")
- s.Require().Empty(groups, "expected 0 groups after remove")
- s.Require().NoError(s.repo.BindGroups(s.ctx, account.ID, []int64{g1.ID, g2.ID}), "BindGroups")
- groups, err = s.repo.GetGroups(s.ctx, account.ID)
- s.Require().NoError(err, "GetGroups after bind")
- s.Require().Len(groups, 2, "expected 2 groups after bind")
- }
- func (s *AccountRepoSuite) TestBindGroups_EmptyList() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-empty"})
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-empty"})
- mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
- s.Require().NoError(s.repo.BindGroups(s.ctx, account.ID, []int64{}), "BindGroups empty")
- groups, err := s.repo.GetGroups(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Empty(groups, "expected 0 groups after binding empty list")
- }
- // --- Schedulable ---
- func (s *AccountRepoSuite) TestListSchedulable() {
- now := time.Now()
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sched"})
- okAcc := mustCreateAccount(s.T(), s.client, &service.Account{Name: "ok", Schedulable: true})
- mustBindAccountToGroup(s.T(), s.client, okAcc.ID, group.ID, 1)
- future := now.Add(10 * time.Minute)
- overloaded := mustCreateAccount(s.T(), s.client, &service.Account{Name: "over", Schedulable: true, OverloadUntil: &future})
- mustBindAccountToGroup(s.T(), s.client, overloaded.ID, group.ID, 1)
- sched, err := s.repo.ListSchedulable(s.ctx)
- s.Require().NoError(err, "ListSchedulable")
- ids := idsOfAccounts(sched)
- s.Require().Contains(ids, okAcc.ID)
- s.Require().NotContains(ids, overloaded.ID)
- }
- func (s *AccountRepoSuite) TestListSchedulableByGroupID_TimeBoundaries_And_StatusUpdates() {
- now := time.Now()
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sched"})
- okAcc := mustCreateAccount(s.T(), s.client, &service.Account{Name: "ok", Schedulable: true})
- mustBindAccountToGroup(s.T(), s.client, okAcc.ID, group.ID, 1)
- future := now.Add(10 * time.Minute)
- overloaded := mustCreateAccount(s.T(), s.client, &service.Account{Name: "over", Schedulable: true, OverloadUntil: &future})
- mustBindAccountToGroup(s.T(), s.client, overloaded.ID, group.ID, 1)
- rateLimited := mustCreateAccount(s.T(), s.client, &service.Account{Name: "rl", Schedulable: true})
- mustBindAccountToGroup(s.T(), s.client, rateLimited.ID, group.ID, 1)
- s.Require().NoError(s.repo.SetRateLimited(s.ctx, rateLimited.ID, now.Add(10*time.Minute)), "SetRateLimited")
- s.Require().NoError(s.repo.SetError(s.ctx, overloaded.ID, "boom"), "SetError")
- sched, err := s.repo.ListSchedulableByGroupID(s.ctx, group.ID)
- s.Require().NoError(err, "ListSchedulableByGroupID")
- s.Require().Len(sched, 1, "expected only ok account schedulable")
- s.Require().Equal(okAcc.ID, sched[0].ID)
- s.Require().NoError(s.repo.ClearRateLimit(s.ctx, rateLimited.ID), "ClearRateLimit")
- sched2, err := s.repo.ListSchedulableByGroupID(s.ctx, group.ID)
- s.Require().NoError(err, "ListSchedulableByGroupID after ClearRateLimit")
- s.Require().Len(sched2, 2, "expected 2 schedulable accounts after ClearRateLimit")
- }
- func (s *AccountRepoSuite) TestListSchedulableByPlatform() {
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic, Schedulable: true})
- mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI, Schedulable: true})
- accounts, err := s.repo.ListSchedulableByPlatform(s.ctx, service.PlatformAnthropic)
- s.Require().NoError(err)
- s.Require().Len(accounts, 1)
- s.Require().Equal(service.PlatformAnthropic, accounts[0].Platform)
- }
- func (s *AccountRepoSuite) TestListSchedulableByGroupIDAndPlatform() {
- group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sp"})
- a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic, Schedulable: true})
- a2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI, Schedulable: true})
- mustBindAccountToGroup(s.T(), s.client, a1.ID, group.ID, 1)
- mustBindAccountToGroup(s.T(), s.client, a2.ID, group.ID, 2)
- accounts, err := s.repo.ListSchedulableByGroupIDAndPlatform(s.ctx, group.ID, service.PlatformAnthropic)
- s.Require().NoError(err)
- s.Require().Len(accounts, 1)
- s.Require().Equal(a1.ID, accounts[0].ID)
- }
- func (s *AccountRepoSuite) TestSetSchedulable() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-sched", Schedulable: true})
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.SetSchedulable(s.ctx, account.ID, false))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().False(got.Schedulable)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- }
- func (s *AccountRepoSuite) TestBulkUpdate_SyncSchedulerSnapshotOnDisabled() {
- account1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-1", Status: service.StatusActive, Schedulable: true})
- account2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-2", Status: service.StatusActive, Schedulable: true})
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- disabled := service.StatusDisabled
- rows, err := s.repo.BulkUpdate(s.ctx, []int64{account1.ID, account2.ID}, service.AccountBulkUpdate{
- Status: &disabled,
- })
- s.Require().NoError(err)
- s.Require().Equal(int64(2), rows)
- s.Require().Len(cacheRecorder.setAccounts, 2)
- ids := map[int64]struct{}{}
- for _, acc := range cacheRecorder.setAccounts {
- ids[acc.ID] = struct{}{}
- }
- s.Require().Contains(ids, account1.ID)
- s.Require().Contains(ids, account2.ID)
- }
- // --- SetOverloaded / SetRateLimited / ClearRateLimit ---
- func (s *AccountRepoSuite) TestSetOverloaded() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-over"})
- until := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(got.OverloadUntil)
- s.Require().WithinDuration(until, *got.OverloadUntil, time.Second)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().NotNil(cacheRecorder.setAccounts[0].OverloadUntil)
- s.Require().WithinDuration(until, *cacheRecorder.setAccounts[0].OverloadUntil, time.Second)
- }
- func (s *AccountRepoSuite) TestSetRateLimited() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl"})
- resetAt := time.Date(2025, 6, 15, 14, 0, 0, 0, time.UTC)
- s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, resetAt))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(got.RateLimitedAt)
- s.Require().NotNil(got.RateLimitResetAt)
- s.Require().WithinDuration(resetAt, *got.RateLimitResetAt, time.Second)
- }
- func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl-monotonic"})
- later := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
- earlier := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second)
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, later))
- s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, earlier))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(got.RateLimitResetAt)
- s.Require().WithinDuration(later, *got.RateLimitResetAt, time.Second)
- s.Require().Len(cacheRecorder.setAccounts, 2)
- s.Require().NotNil(cacheRecorder.setAccounts[1].RateLimitResetAt)
- s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second)
- }
- func (s *AccountRepoSuite) TestClearRateLimitIfObservedProtectsRearmed429Generation() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-rl-conditional-clear",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- })
- firstReset := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
- rearmedReset := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second)
- s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, firstReset))
- staleGeneration, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(staleGeneration.RateLimitedAt)
- s.Require().NotNil(staleGeneration.RateLimitResetAt)
- cleared, err := s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt)
- s.Require().NoError(err)
- s.Require().True(cleared)
- // A newer generation may legitimately re-arm a shorter boundary after the
- // first generation was cleared. The stale success must not erase it.
- s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, rearmedReset))
- cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt)
- s.Require().NoError(err)
- s.Require().False(cleared)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(got.RateLimitedAt)
- s.Require().NotNil(got.RateLimitResetAt)
- s.Require().WithinDuration(rearmedReset, *got.RateLimitResetAt, time.Second)
- // An admin can retype the row while the successful OAuth request is still
- // in flight. The stale OAuth recovery must not cross into API-key state even
- // when both observed timestamps still match.
- _, err = s.client.Account.UpdateOneID(account.ID).
- SetType(service.AccountTypeAPIKey).
- Save(s.ctx)
- s.Require().NoError(err)
- cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *got.RateLimitedAt, *got.RateLimitResetAt)
- s.Require().NoError(err)
- s.Require().False(cleared)
- retyped, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal(service.AccountTypeAPIKey, retyped.Type)
- s.Require().NotNil(retyped.RateLimitedAt)
- s.Require().NotNil(retyped.RateLimitResetAt)
- s.Require().WithinDuration(rearmedReset, *retyped.RateLimitResetAt, time.Second)
- }
- func (s *AccountRepoSuite) TestClearRateLimit() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"})
- until := time.Now().Add(1 * time.Hour)
- s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
- s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, until))
- s.Require().NoError(s.repo.ClearRateLimit(s.ctx, account.ID))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Nil(got.RateLimitedAt)
- s.Require().Nil(got.RateLimitResetAt)
- s.Require().Nil(got.OverloadUntil)
- }
- func (s *AccountRepoSuite) TestResetQuotaUsedAndClearRateLimitCooldownPreservesOtherRuntimeState() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-reset-quota-cooldown",
- Extra: map[string]any{
- "quota_used": 12.5,
- "quota_daily_used": 5.0,
- "quota_weekly_used": 9.0,
- "model_rate_limits": map[string]any{
- "claude-sonnet-4-5": map[string]any{"rate_limit_reset_at": "2026-09-01T10:00:00Z"},
- },
- },
- })
- until := time.Now().Add(1 * time.Hour)
- s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
- s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, until))
- s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, account.ID, until, "preserve-me"))
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.ResetQuotaUsedAndClearRateLimitCooldown(s.ctx, account.ID))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Nil(got.RateLimitedAt)
- s.Require().Nil(got.RateLimitResetAt)
- s.Require().NotNil(got.OverloadUntil)
- s.Require().WithinDuration(until, *got.OverloadUntil, time.Second)
- s.Require().NotNil(got.TempUnschedulableUntil)
- s.Require().WithinDuration(until, *got.TempUnschedulableUntil, time.Second)
- s.Require().Equal("preserve-me", got.TempUnschedulableReason)
- s.Require().Contains(got.Extra, "model_rate_limits")
- s.Require().Equal(float64(0), got.Extra["quota_used"])
- s.Require().Equal(float64(0), got.Extra["quota_daily_used"])
- s.Require().Equal(float64(0), got.Extra["quota_weekly_used"])
- var pendingEventExists bool
- s.Require().NoError(scanSingleRow(s.ctx, s.repo.sql, `
- SELECT EXISTS (
- SELECT 1 FROM scheduler_outbox
- WHERE event_type = $1 AND account_id = $2 AND dedup_key IS NOT NULL
- )`, []any{service.SchedulerOutboxEventAccountChanged, account.ID}, &pendingEventExists))
- s.Require().True(pendingEventExists)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- }
- func (s *AccountRepoSuite) TestTempUnschedulableFieldsLoadedByGetByIDAndGetByIDs() {
- acc1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-temp-1"})
- acc2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-temp-2"})
- until := time.Now().Add(15 * time.Minute).UTC().Truncate(time.Second)
- reason := `{"rule":"429","matched_keyword":"too many requests"}`
- s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, acc1.ID, until, reason))
- gotByID, err := s.repo.GetByID(s.ctx, acc1.ID)
- s.Require().NoError(err)
- s.Require().NotNil(gotByID.TempUnschedulableUntil)
- s.Require().WithinDuration(until, *gotByID.TempUnschedulableUntil, time.Second)
- s.Require().Equal(reason, gotByID.TempUnschedulableReason)
- gotByIDs, err := s.repo.GetByIDs(s.ctx, []int64{acc2.ID, acc1.ID})
- s.Require().NoError(err)
- s.Require().Len(gotByIDs, 2)
- s.Require().Equal(acc2.ID, gotByIDs[0].ID)
- s.Require().Nil(gotByIDs[0].TempUnschedulableUntil)
- s.Require().Equal("", gotByIDs[0].TempUnschedulableReason)
- s.Require().Equal(acc1.ID, gotByIDs[1].ID)
- s.Require().NotNil(gotByIDs[1].TempUnschedulableUntil)
- s.Require().WithinDuration(until, *gotByIDs[1].TempUnschedulableUntil, time.Second)
- s.Require().Equal(reason, gotByIDs[1].TempUnschedulableReason)
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.ClearTempUnschedulable(s.ctx, acc1.ID))
- cleared, err := s.repo.GetByID(s.ctx, acc1.ID)
- s.Require().NoError(err)
- s.Require().Nil(cleared.TempUnschedulableUntil)
- s.Require().Equal("", cleared.TempUnschedulableReason)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(acc1.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().Nil(cacheRecorder.setAccounts[0].TempUnschedulableUntil)
- s.Require().Equal("", cacheRecorder.setAccounts[0].TempUnschedulableReason)
- }
- func (s *AccountRepoSuite) TestSetTempUnschedulableSkipsOutboxWhenWindowDoesNotExtend() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-temp-noop"})
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- until := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
- s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, account.ID, until, "first"))
- var count int
- err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
- s.Require().NoError(err)
- s.Require().Equal(1, count)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().NoError(s.repo.SetTempUnschedulable(s.ctx, account.ID, until.Add(-5*time.Minute), "older"))
- err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
- s.Require().NoError(err)
- s.Require().Equal(1, count)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal("first", got.TempUnschedulableReason)
- s.Require().NotNil(got.TempUnschedulableUntil)
- s.Require().WithinDuration(until, *got.TempUnschedulableUntil, time.Second)
- }
- func (s *AccountRepoSuite) TestClearModelRateLimits_SyncsSchedulerSnapshot() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-clear-model-rate",
- Extra: map[string]any{
- "model_rate_limits": map[string]any{
- "claude-sonnet-4-5": map[string]any{
- "rate_limit_reset_at": "2026-06-03T10:00:00Z",
- },
- },
- },
- })
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.ClearModelRateLimits(s.ctx, account.ID))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotContains(got.Extra, "model_rate_limits")
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().NotContains(cacheRecorder.setAccounts[0].Extra, "model_rate_limits")
- }
- // --- UpdateLastUsed ---
- func (s *AccountRepoSuite) TestUpdateLastUsed() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-used"})
- s.Require().Nil(account.LastUsedAt)
- s.Require().NoError(s.repo.UpdateLastUsed(s.ctx, account.ID))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(got.LastUsedAt)
- }
- // --- SetError ---
- func (s *AccountRepoSuite) TestSetError() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-err", Status: service.StatusActive, Schedulable: true})
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- s.Require().NoError(s.repo.SetError(s.ctx, account.ID, "something went wrong"))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal(service.StatusError, got.Status)
- s.Require().Equal("something went wrong", got.ErrorMessage)
- s.Require().False(got.Schedulable)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().Equal(service.StatusError, cacheRecorder.setAccounts[0].Status)
- s.Require().False(cacheRecorder.setAccounts[0].Schedulable)
- var outboxCount int
- err = scanSingleRow(
- s.ctx,
- s.repo.sql,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- []any{service.SchedulerOutboxEventAccountChanged, account.ID},
- &outboxCount,
- )
- s.Require().NoError(err)
- s.Require().Equal(1, outboxCount)
- }
- func (s *AccountRepoSuite) TestSetGrokOAuthErrorIfCredentialsUnchanged_AppliesAndSyncsSchedulerState() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-conditional-error-applied",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{"access_token": "observed", "_token_version": int64(7)},
- })
- observed, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- applied, err := s.repo.SetGrokOAuthErrorIfCredentialsUnchanged(
- s.ctx,
- account.ID,
- observed.Credentials,
- "missing refresh token",
- )
- s.Require().NoError(err)
- s.Require().True(applied)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal(service.StatusError, got.Status)
- s.Require().False(got.Schedulable)
- s.Require().Equal("missing refresh token", got.ErrorMessage)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(service.StatusError, cacheRecorder.setAccounts[0].Status)
- var outboxCount int
- err = scanSingleRow(
- s.ctx,
- s.repo.sql,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- []any{service.SchedulerOutboxEventAccountChanged, account.ID},
- &outboxCount,
- )
- s.Require().NoError(err)
- s.Require().Equal(1, outboxCount)
- }
- func (s *AccountRepoSuite) TestSetGrokOAuthErrorIfCredentialsUnchanged_SkipsConcurrentReauthorization() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-conditional-error-reauthorized",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{"access_token": "observed", "_token_version": int64(7)},
- })
- observed, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{
- "access_token": "fresh-access",
- "refresh_token": "fresh-refresh",
- "expires_at": time.Now().UTC().Add(4 * time.Hour).Format(time.RFC3339),
- "_token_version": int64(8),
- }))
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- applied, err := s.repo.SetGrokOAuthErrorIfCredentialsUnchanged(
- s.ctx,
- account.ID,
- observed.Credentials,
- "stale reconciliation",
- )
- s.Require().NoError(err)
- s.Require().False(applied)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal(service.StatusActive, got.Status)
- s.Require().True(got.Schedulable)
- s.Require().Equal("fresh-refresh", got.GetGrokRefreshToken())
- s.Require().Empty(cacheRecorder.setAccounts, "a lost compare-and-set race must not rewrite the scheduler snapshot")
- var outboxCount int
- err = scanSingleRow(
- s.ctx,
- s.repo.sql,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- []any{service.SchedulerOutboxEventAccountChanged, account.ID},
- &outboxCount,
- )
- s.Require().NoError(err)
- s.Require().Zero(outboxCount, "a lost compare-and-set race must not enqueue a stale account change")
- }
- func (s *AccountRepoSuite) TestUpdateGrokOAuthCredentialsIfUnchanged_AppliesAndPublishesSchedulerState() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-refresh-success-cas-applied",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{
- "access_token": "attempted-access",
- "refresh_token": "attempted-refresh",
- "_token_version": int64(10),
- },
- })
- observed, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- applied, err := s.repo.UpdateGrokOAuthCredentialsIfUnchanged(
- s.ctx,
- account.ID,
- observed.Credentials,
- observed.ProxyID,
- map[string]any{
- "access_token": "rotated-access",
- "refresh_token": "rotated-refresh",
- "_token_version": int64(11),
- },
- )
- s.Require().NoError(err)
- s.Require().True(applied)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal("rotated-refresh", got.GetGrokRefreshToken())
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal("rotated-refresh", cacheRecorder.setAccounts[0].GetGrokRefreshToken())
- s.Require().NoError(cacheRecorder.setCtxErr)
- var outboxCount int
- err = scanSingleRow(
- s.ctx,
- s.repo.sql,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- []any{service.SchedulerOutboxEventAccountChanged, account.ID},
- &outboxCount,
- )
- s.Require().NoError(err)
- s.Require().Equal(1, outboxCount)
- }
- func (s *AccountRepoSuite) TestUpdateGrokOAuthCredentialsIfUnchanged_SkipsConcurrentReauthorization() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-refresh-success-cas-reauthorized",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{
- "access_token": "attempted-access",
- "refresh_token": "attempted-refresh",
- "_token_version": int64(20),
- },
- })
- observed, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NoError(s.repo.UpdateCredentials(s.ctx, account.ID, map[string]any{
- "access_token": "reauthorized-access",
- "refresh_token": "reauthorized-refresh",
- "_token_version": int64(21),
- }))
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err = s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- applied, err := s.repo.UpdateGrokOAuthCredentialsIfUnchanged(
- s.ctx,
- account.ID,
- observed.Credentials,
- observed.ProxyID,
- map[string]any{
- "access_token": "provider-access",
- "refresh_token": "provider-refresh",
- "_token_version": int64(22),
- },
- )
- s.Require().NoError(err)
- s.Require().False(applied)
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal("reauthorized-refresh", got.GetGrokRefreshToken())
- s.Require().Empty(cacheRecorder.setAccounts)
- var outboxCount int
- err = scanSingleRow(
- s.ctx,
- s.repo.sql,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- []any{service.SchedulerOutboxEventAccountChanged, account.ID},
- &outboxCount,
- )
- s.Require().NoError(err)
- s.Require().Zero(outboxCount)
- }
- func (s *AccountRepoSuite) TestGrokOAuthConditionalMutation_DetachesBoundedSnapshotSync() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "grok-conditional-detached-sync",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{"access_token": "observed"},
- })
- observed, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- ctx, cancel := context.WithCancel(context.Background())
- cacheRecorder := &schedulerCacheRecorder{}
- repo := newAccountRepositoryWithSQL(s.client, &cancelAfterAtomicMutationSQLExecutor{
- sqlExecutor: s.repo.sql,
- cancel: cancel,
- }, cacheRecorder)
- applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
- ctx,
- account.ID,
- observed.Credentials,
- "missing refresh token",
- )
- s.Require().NoError(err)
- s.Require().True(applied)
- s.Require().ErrorIs(ctx.Err(), context.Canceled)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().NoError(cacheRecorder.setCtxErr, "immediate scheduler propagation must use a bounded detached context")
- }
- func TestGrokOAuthConditionalMutationRollsBackWhenOutboxInsertFails(t *testing.T) {
- client := testEntClient(t)
- account := mustCreateAccount(t, client, &service.Account{
- Name: "grok-conditional-atomic-outbox-failure",
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Status: service.StatusActive,
- Schedulable: true,
- Credentials: map[string]any{"access_token": "observed"},
- })
- t.Cleanup(func() {
- _, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM scheduler_outbox WHERE account_id = $1", account.ID)
- _ = client.Account.DeleteOneID(account.ID).Exec(context.Background())
- })
- repo := newAccountRepositoryWithSQL(client, &failAtomicSchedulerOutboxSQLExecutor{sqlExecutor: integrationDB}, nil)
- applied, err := repo.SetGrokOAuthErrorIfCredentialsUnchanged(
- context.Background(),
- account.ID,
- account.Credentials,
- "missing refresh token",
- )
- require.Error(t, err)
- require.False(t, applied)
- got, readErr := repo.GetByID(context.Background(), account.ID)
- require.NoError(t, readErr)
- require.Equal(t, service.StatusActive, got.Status)
- require.True(t, got.Schedulable)
- require.Empty(t, got.ErrorMessage)
- var outboxCount int
- require.NoError(t, integrationDB.QueryRowContext(
- context.Background(),
- "SELECT COUNT(*) FROM scheduler_outbox WHERE account_id = $1",
- account.ID,
- ).Scan(&outboxCount))
- require.Zero(t, outboxCount)
- }
- func (s *AccountRepoSuite) TestUpdateErrorStatusUnschedulesAccount() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-update-err", Status: service.StatusActive, Schedulable: true})
- account.Status = service.StatusError
- account.ErrorMessage = "token revoked"
- account.Schedulable = true
- s.Require().NoError(s.repo.Update(s.ctx, account))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal(service.StatusError, got.Status)
- s.Require().Equal("token revoked", got.ErrorMessage)
- s.Require().False(got.Schedulable)
- }
- func (s *AccountRepoSuite) TestClearError_SyncSchedulerSnapshotOnRecovery() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-clear-err",
- Status: service.StatusError,
- ErrorMessage: "temporary error",
- })
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- s.Require().NoError(s.repo.ClearError(s.ctx, account.ID))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal(service.StatusActive, got.Status)
- s.Require().Empty(got.ErrorMessage)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().Equal(service.StatusActive, cacheRecorder.setAccounts[0].Status)
- }
- // --- UpdateSessionWindow ---
- func (s *AccountRepoSuite) TestUpdateSessionWindow() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-win"})
- start := time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC)
- end := time.Date(2025, 6, 15, 15, 0, 0, 0, time.UTC)
- s.Require().NoError(s.repo.UpdateSessionWindow(s.ctx, account.ID, &start, &end, "active"))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().NotNil(got.SessionWindowStart)
- s.Require().NotNil(got.SessionWindowEnd)
- s.Require().Equal("active", got.SessionWindowStatus)
- }
- // --- UpdateExtra ---
- func (s *AccountRepoSuite) TestUpdateExtra_MergesFields() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-extra",
- Extra: map[string]any{"a": "1"},
- })
- s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{"b": "2"}), "UpdateExtra")
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err, "GetByID")
- s.Require().Equal("1", got.Extra["a"])
- s.Require().Equal("2", got.Extra["b"])
- }
- func (s *AccountRepoSuite) TestUpdateExtra_EmptyUpdates() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-extra-empty"})
- s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{}))
- }
- func (s *AccountRepoSuite) TestUpdateExtra_NilExtra() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-nil-extra", Extra: nil})
- s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{"key": "val"}))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal("val", got.Extra["key"])
- }
- func (s *AccountRepoSuite) TestUpdateExtra_SchedulerNeutralSkipsOutboxAndSyncsFreshSnapshot() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-extra-neutral",
- Platform: service.PlatformOpenAI,
- Extra: map[string]any{"codex_usage_updated_at": "old"},
- })
- cacheRecorder := &schedulerCacheRecorder{
- accounts: map[int64]*service.Account{
- account.ID: {
- ID: account.ID,
- Platform: account.Platform,
- Status: service.StatusDisabled,
- Extra: map[string]any{
- "codex_usage_updated_at": "old",
- },
- },
- },
- }
- s.repo.schedulerCache = cacheRecorder
- updates := map[string]any{
- "codex_usage_updated_at": "2026-03-11T10:00:00Z",
- "codex_5h_used_percent": 88.5,
- "session_window_utilization": 0.42,
- }
- s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, updates))
- got, err := s.repo.GetByID(s.ctx, account.ID)
- s.Require().NoError(err)
- s.Require().Equal("2026-03-11T10:00:00Z", got.Extra["codex_usage_updated_at"])
- s.Require().Equal(88.5, got.Extra["codex_5h_used_percent"])
- s.Require().Equal(0.42, got.Extra["session_window_utilization"])
- var outboxCount int
- s.Require().NoError(scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &outboxCount))
- s.Require().Zero(outboxCount)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().NotNil(cacheRecorder.accounts[account.ID])
- s.Require().Equal(service.StatusActive, cacheRecorder.accounts[account.ID].Status)
- s.Require().Equal("2026-03-11T10:00:00Z", cacheRecorder.accounts[account.ID].Extra["codex_usage_updated_at"])
- }
- func (s *AccountRepoSuite) TestUpdateExtra_ExhaustedCodexSnapshotSyncsSchedulerCache() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-extra-codex-exhausted",
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Extra: map[string]any{},
- })
- cacheRecorder := &schedulerCacheRecorder{}
- s.repo.schedulerCache = cacheRecorder
- _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{
- "codex_7d_used_percent": 100.0,
- "codex_7d_reset_at": "2026-03-12T13:00:00Z",
- "codex_7d_reset_after_seconds": 86400,
- }))
- var count int
- err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
- s.Require().NoError(err)
- s.Require().Equal(0, count)
- s.Require().Len(cacheRecorder.setAccounts, 1)
- s.Require().Equal(account.ID, cacheRecorder.setAccounts[0].ID)
- s.Require().Equal(service.StatusActive, cacheRecorder.setAccounts[0].Status)
- s.Require().Equal(100.0, cacheRecorder.setAccounts[0].Extra["codex_7d_used_percent"])
- }
- func (s *AccountRepoSuite) TestUpdateExtra_SchedulerRelevantStillEnqueuesOutbox() {
- account := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-extra-mixed",
- Platform: service.PlatformAntigravity,
- Extra: map[string]any{},
- })
- _, err := s.repo.sql.ExecContext(s.ctx, "TRUNCATE scheduler_outbox")
- s.Require().NoError(err)
- s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{
- "mixed_scheduling": true,
- "codex_usage_updated_at": "2026-03-11T10:00:00Z",
- }))
- var count int
- err = scanSingleRow(s.ctx, s.repo.sql, "SELECT COUNT(*) FROM scheduler_outbox", nil, &count)
- s.Require().NoError(err)
- s.Require().Equal(1, count)
- }
- // --- GetByCRSAccountID ---
- func (s *AccountRepoSuite) TestGetByCRSAccountID() {
- crsID := "crs-12345"
- mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "acc-crs",
- Extra: map[string]any{"crs_account_id": crsID},
- })
- got, err := s.repo.GetByCRSAccountID(s.ctx, crsID)
- s.Require().NoError(err)
- s.Require().NotNil(got)
- s.Require().Equal("acc-crs", got.Name)
- }
- func (s *AccountRepoSuite) TestGetByCRSAccountID_NotFound() {
- got, err := s.repo.GetByCRSAccountID(s.ctx, "non-existent")
- s.Require().NoError(err)
- s.Require().Nil(got)
- }
- func (s *AccountRepoSuite) TestGetByCRSAccountID_EmptyString() {
- got, err := s.repo.GetByCRSAccountID(s.ctx, "")
- s.Require().NoError(err)
- s.Require().Nil(got)
- }
- // TestGetByCRSAccountID_ExcludesSparkShadow 验证外审第7轮 P1:即便 spark 影子的 Extra 被误写入
- // crs_account_id,CRS 查询也绝不能命中影子(否则会被当普通账号更新而覆盖 type/credentials/proxy)。
- func (s *AccountRepoSuite) TestGetByCRSAccountID_ExcludesSparkShadow() {
- crsID := "crs-shadow-only-99"
- parent := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "crs-mother", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
- })
- mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "crs-shadow", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
- ParentAccountID: &parent.ID,
- QuotaDimension: service.QuotaDimensionSpark,
- Extra: map[string]any{"crs_account_id": crsID},
- })
- got, err := s.repo.GetByCRSAccountID(s.ctx, crsID)
- s.Require().NoError(err)
- s.Require().Nil(got, "spark 影子即便带 crs_account_id 也不应被 CRS 命中")
- }
- // TestListCRSAccountIDs_ExcludesSparkShadow 验证外审第7轮 P1:影子的 crs_account_id 不应进入
- // CRS 同步映射(否则后续 CRS 同步会把影子当普通账号更新)。
- func (s *AccountRepoSuite) TestListCRSAccountIDs_ExcludesSparkShadow() {
- parent := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "crs-list-mother", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
- })
- shadowCRSID := "crs-list-shadow-77"
- mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "crs-list-shadow", Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
- ParentAccountID: &parent.ID,
- QuotaDimension: service.QuotaDimensionSpark,
- Extra: map[string]any{"crs_account_id": shadowCRSID},
- })
- ids, err := s.repo.ListCRSAccountIDs(s.ctx)
- s.Require().NoError(err)
- _, ok := ids[shadowCRSID]
- s.Require().False(ok, "影子的 crs_account_id 不应进入 CRS 映射")
- }
- // --- BulkUpdate ---
- func (s *AccountRepoSuite) TestBulkUpdate() {
- a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk1", Priority: 1})
- a2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk2", Priority: 1})
- newPriority := 99
- affected, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID, a2.ID}, service.AccountBulkUpdate{
- Priority: &newPriority,
- })
- s.Require().NoError(err)
- s.Require().GreaterOrEqual(affected, int64(1), "expected at least one affected row")
- got1, _ := s.repo.GetByID(s.ctx, a1.ID)
- got2, _ := s.repo.GetByID(s.ctx, a2.ID)
- s.Require().Equal(99, got1.Priority)
- s.Require().Equal(99, got2.Priority)
- }
- func (s *AccountRepoSuite) TestBulkUpdate_MergeCredentials() {
- a1 := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "bulk-cred",
- Credentials: map[string]any{"existing": "value"},
- })
- _, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{
- Credentials: map[string]any{"new_key": "new_value"},
- })
- s.Require().NoError(err)
- got, _ := s.repo.GetByID(s.ctx, a1.ID)
- s.Require().Equal("value", got.Credentials["existing"])
- s.Require().Equal("new_value", got.Credentials["new_key"])
- }
- func (s *AccountRepoSuite) TestBulkUpdate_MergeExtra() {
- a1 := mustCreateAccount(s.T(), s.client, &service.Account{
- Name: "bulk-extra",
- Extra: map[string]any{"existing": "val"},
- })
- _, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{
- Extra: map[string]any{"new_key": "new_val"},
- })
- s.Require().NoError(err)
- got, _ := s.repo.GetByID(s.ctx, a1.ID)
- s.Require().Equal("val", got.Extra["existing"])
- s.Require().Equal("new_val", got.Extra["new_key"])
- }
- func (s *AccountRepoSuite) TestBulkUpdate_EmptyIDs() {
- affected, err := s.repo.BulkUpdate(s.ctx, []int64{}, service.AccountBulkUpdate{})
- s.Require().NoError(err)
- s.Require().Zero(affected)
- }
- func (s *AccountRepoSuite) TestBulkUpdate_EmptyUpdates() {
- a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-empty"})
- affected, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{})
- s.Require().NoError(err)
- s.Require().Zero(affected)
- }
- func idsOfAccounts(accounts []service.Account) []int64 {
- out := make([]int64, 0, len(accounts))
- for i := range accounts {
- out = append(out, accounts[i].ID)
- }
- return out
- }
|