//go:build integration package repository import ( "context" "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/suite" ) type AccountRepoSuite struct { suite.Suite ctx context.Context client *dbent.Client repo *accountRepository } func (s *AccountRepoSuite) SetupTest() { s.ctx = context.Background() tx := testEntTx(s.T()) s.client = tx.Client() s.repo = newAccountRepositoryWithSQL(s.client, tx) } 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) 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_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) TestListWithFilters() { tests := []struct { name string setup func(client *dbent.Client) platform string accType string status string search 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_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") }, }, } for _, tt := range tests { s.Run(tt.name, func() { // 每个 case 重新获取隔离资源 tx := testEntTx(s.T()) client := tx.Client() repo := newAccountRepositoryWithSQL(client, tx) ctx := context.Background() tt.setup(client) accounts, _, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, tt.platform, tt.accType, tt.status, tt.search) s.Require().NoError(err) s.Require().Len(accounts, tt.wantCount) 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") 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}) 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) } // --- 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) 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) } 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) 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) } // --- 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}) 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) } // --- 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"]) } // --- 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) } // --- 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 }