admin_service_duplicate_account_test.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "net/http"
  7. "strings"
  8. "testing"
  9. "time"
  10. "unicode/utf8"
  11. "github.com/Wei-Shaw/sub2api/internal/config"
  12. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  13. "github.com/stretchr/testify/require"
  14. )
  15. type duplicateAccountRepoStub struct {
  16. *sparkShadowRepoStub
  17. atomicCreateErr error
  18. accountGroupsOf map[int64][]AccountGroup
  19. }
  20. func newDuplicateAccountRepoStub() *duplicateAccountRepoStub {
  21. return &duplicateAccountRepoStub{
  22. sparkShadowRepoStub: newSparkShadowRepoStub(),
  23. accountGroupsOf: make(map[int64][]AccountGroup),
  24. }
  25. }
  26. func (s *duplicateAccountRepoStub) CreateWithAccountGroups(ctx context.Context, account *Account, groups []AccountGroup) error {
  27. if s.atomicCreateErr != nil {
  28. return s.atomicCreateErr
  29. }
  30. groupIDs := make([]int64, 0, len(groups))
  31. for _, group := range groups {
  32. groupIDs = append(groupIDs, group.GroupID)
  33. }
  34. account.GroupIDs = groupIDs
  35. if err := s.Create(ctx, account); err != nil {
  36. return err
  37. }
  38. clonedGroups := make([]AccountGroup, len(groups))
  39. copy(clonedGroups, groups)
  40. for i := range clonedGroups {
  41. clonedGroups[i].AccountID = account.ID
  42. }
  43. account.AccountGroups = clonedGroups
  44. s.accountGroupsOf[account.ID] = clonedGroups
  45. if len(groupIDs) > 0 {
  46. s.groupsOf[account.ID] = append([]int64(nil), groupIDs...)
  47. }
  48. stored := *account
  49. s.accounts[account.ID] = &stored
  50. s.mockAccountRepoForGemini.accountsByID[account.ID] = &stored
  51. return nil
  52. }
  53. func (s *duplicateAccountRepoStub) FindByExtraField(_ context.Context, key string, value any) ([]Account, error) {
  54. wanted, ok := value.(string)
  55. if !ok {
  56. return nil, nil
  57. }
  58. var matches []Account
  59. for _, account := range s.accounts {
  60. if actual, ok := account.Extra[key].(string); ok && actual == wanted {
  61. matches = append(matches, *account)
  62. }
  63. }
  64. return matches, nil
  65. }
  66. func TestDuplicateAccountCopiesConfigurationAndResetsRuntimeState(t *testing.T) {
  67. ctx := context.Background()
  68. repo := newDuplicateAccountRepoStub()
  69. svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
  70. notes := "keep this note"
  71. proxyID := int64(17)
  72. originalProxyID := int64(11)
  73. rateMultiplier := 1.25
  74. loadFactor := 9
  75. expiresAt := time.Date(2027, time.March, 4, 5, 6, 7, 0, time.UTC)
  76. rateLimitedAt := time.Now().Add(-time.Minute)
  77. rateLimitResetAt := time.Now().Add(time.Hour)
  78. overloadUntil := time.Now().Add(2 * time.Hour)
  79. tempUnschedulableUntil := time.Now().Add(3 * time.Hour)
  80. sessionWindowStart := time.Now().Add(-2 * time.Hour)
  81. sessionWindowEnd := time.Now().Add(2 * time.Hour)
  82. source := &Account{
  83. Name: "primary",
  84. Notes: &notes,
  85. Platform: PlatformAnthropic,
  86. Type: AccountTypeAPIKey,
  87. ProxyID: &proxyID,
  88. ProxyFallbackOriginID: &originalProxyID,
  89. Concurrency: 6,
  90. Priority: 40,
  91. RateMultiplier: &rateMultiplier,
  92. LoadFactor: &loadFactor,
  93. Status: StatusError,
  94. Schedulable: true,
  95. ErrorMessage: "upstream unavailable",
  96. ExpiresAt: &expiresAt,
  97. AutoPauseOnExpired: false,
  98. Credentials: map[string]any{
  99. "api_key": "secret",
  100. "nested": map[string]any{"token": "source-token"},
  101. },
  102. Extra: map[string]any{
  103. "config": map[string]any{"region": "us-east-1"},
  104. "items": []any{map[string]any{"enabled": true}},
  105. "quota_limit": 1000,
  106. "quota_used": 450,
  107. "quota_daily_used": 25,
  108. "quota_daily_start": "2026-07-15T00:00:00Z",
  109. "model_rate_limits": map[string]any{"gpt-5": "2099-01-01T00:00:00Z"},
  110. "codex_5h_used_percent": 80,
  111. "codex_cli_only": true,
  112. "grok_usage_snapshot": map[string]any{"status_code": 429},
  113. "openai_responses_supported": false,
  114. "openai_compact_checked_at": "2026-07-15T00:00:00Z",
  115. "session_window_utilization": 0.8,
  116. "passive_usage_sampled_at": "2026-07-15T00:00:00Z",
  117. "antigravity_force_token_refresh": true,
  118. "antigravity_credits_overages": map[string]any{"enabled": true},
  119. "crs_account_id": "remote-42",
  120. "crs_kind": "openai-api-key",
  121. "crs_synced_at": "2026-07-15T00:00:00Z",
  122. },
  123. GroupIDs: []int64{7, 3},
  124. AccountGroups: []AccountGroup{{GroupID: 7, Priority: 50}, {GroupID: 3, Priority: 7}},
  125. RateLimitedAt: &rateLimitedAt,
  126. RateLimitResetAt: &rateLimitResetAt,
  127. OverloadUntil: &overloadUntil,
  128. TempUnschedulableUntil: &tempUnschedulableUntil,
  129. TempUnschedulableReason: "maintenance",
  130. SessionWindowStart: &sessionWindowStart,
  131. SessionWindowEnd: &sessionWindowEnd,
  132. SessionWindowStatus: "active",
  133. }
  134. source.Extra[UpstreamBillingProbeEnabledExtraKey] = true
  135. source.Extra[UpstreamBillingRateSyncEnabledExtraKey] = true
  136. source.Extra[UpstreamBillingProbeExtraKey] = map[string]any{"status": "ok"}
  137. require.NoError(t, repo.Create(ctx, source))
  138. duplicate, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
  139. require.NoError(t, err)
  140. require.NotEqual(t, source.ID, duplicate.ID)
  141. require.Equal(t, "primary (Copy)", duplicate.Name)
  142. require.Equal(t, source.Platform, duplicate.Platform)
  143. require.Equal(t, source.Type, duplicate.Type)
  144. require.Equal(t, source.Concurrency, duplicate.Concurrency)
  145. require.Equal(t, source.Priority, duplicate.Priority)
  146. require.Equal(t, source.AutoPauseOnExpired, duplicate.AutoPauseOnExpired)
  147. require.Equal(t, source.GroupIDs, duplicate.GroupIDs)
  148. require.Equal(t, source.Credentials, duplicate.Credentials)
  149. require.Equal(t, map[string]any{
  150. "config": map[string]any{"region": "us-east-1"},
  151. "items": []any{map[string]any{"enabled": true}},
  152. "quota_limit": float64(1000),
  153. "codex_cli_only": true,
  154. }, duplicate.Extra)
  155. require.NotContains(t, duplicate.Extra, UpstreamBillingRateSyncEnabledExtraKey)
  156. require.NotNil(t, duplicate.ExpiresAt)
  157. require.True(t, source.ExpiresAt.Equal(*duplicate.ExpiresAt))
  158. require.Equal(t, source.Notes, duplicate.Notes)
  159. require.Equal(t, source.ProxyFallbackOriginID, duplicate.ProxyID)
  160. require.Equal(t, source.RateMultiplier, duplicate.RateMultiplier)
  161. require.Equal(t, source.LoadFactor, duplicate.LoadFactor)
  162. require.Equal(t, source.GroupIDs, repo.groupsOf[duplicate.ID])
  163. require.Equal(t, []AccountGroup{
  164. {AccountID: duplicate.ID, GroupID: 7, Priority: 50},
  165. {AccountID: duplicate.ID, GroupID: 3, Priority: 7},
  166. }, repo.accountGroupsOf[duplicate.ID])
  167. require.Equal(t, StatusActive, duplicate.Status)
  168. require.False(t, duplicate.Schedulable)
  169. require.Empty(t, duplicate.ErrorMessage)
  170. require.Nil(t, duplicate.LastUsedAt)
  171. require.Nil(t, duplicate.RateLimitedAt)
  172. require.Nil(t, duplicate.RateLimitResetAt)
  173. require.Nil(t, duplicate.OverloadUntil)
  174. require.Nil(t, duplicate.TempUnschedulableUntil)
  175. require.Empty(t, duplicate.TempUnschedulableReason)
  176. require.Nil(t, duplicate.SessionWindowStart)
  177. require.Nil(t, duplicate.SessionWindowEnd)
  178. require.Empty(t, duplicate.SessionWindowStatus)
  179. duplicate.Credentials["nested"].(map[string]any)["token"] = "changed"
  180. duplicate.Extra["config"].(map[string]any)["region"] = "changed"
  181. duplicate.Extra["items"].([]any)[0].(map[string]any)["enabled"] = false
  182. storedSource, getErr := repo.GetByID(ctx, source.ID)
  183. require.NoError(t, getErr)
  184. require.Equal(t, "source-token", storedSource.Credentials["nested"].(map[string]any)["token"])
  185. require.Equal(t, "us-east-1", storedSource.Extra["config"].(map[string]any)["region"])
  186. require.Equal(t, true, storedSource.Extra["items"].([]any)[0].(map[string]any)["enabled"])
  187. require.Equal(t, "remote-42", storedSource.Extra["crs_account_id"])
  188. }
  189. func TestDuplicateAccountRejectsCredentialShadow(t *testing.T) {
  190. ctx := context.Background()
  191. repo := newDuplicateAccountRepoStub()
  192. svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
  193. parentID := int64(99)
  194. shadow := &Account{
  195. Name: "shadow",
  196. Platform: PlatformOpenAI,
  197. Type: AccountTypeOAuth,
  198. ParentAccountID: &parentID,
  199. QuotaDimension: QuotaDimensionSpark,
  200. }
  201. require.NoError(t, repo.Create(ctx, shadow))
  202. _, err := svc.DuplicateAccount(ctx, shadow.ID, "admin:1", "")
  203. require.Error(t, err)
  204. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  205. require.Equal(t, "ACCOUNT_DUPLICATE_SHADOW_UNSUPPORTED", infraerrors.Reason(err))
  206. require.Len(t, repo.accounts, 1)
  207. }
  208. func TestDuplicateAccountRejectsRotatingOrUnknownCredentialTypes(t *testing.T) {
  209. for _, accountType := range []string{AccountTypeOAuth, AccountTypeSetupToken, "legacy-cookie"} {
  210. t.Run(accountType, func(t *testing.T) {
  211. ctx := context.Background()
  212. repo := newDuplicateAccountRepoStub()
  213. svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
  214. source := &Account{
  215. Name: "rotating-credential-account",
  216. Platform: PlatformOpenAI,
  217. Type: accountType,
  218. Credentials: map[string]any{"refresh_token": "shared-token"},
  219. }
  220. require.NoError(t, repo.Create(ctx, source))
  221. _, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
  222. require.Error(t, err)
  223. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  224. require.Equal(t, "ACCOUNT_DUPLICATE_CREDENTIAL_TYPE_UNSUPPORTED", infraerrors.Reason(err))
  225. require.Len(t, repo.accounts, 1)
  226. })
  227. }
  228. }
  229. func TestDuplicateAccountPreservesUngroupedState(t *testing.T) {
  230. ctx := context.Background()
  231. repo := newDuplicateAccountRepoStub()
  232. svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
  233. source := &Account{
  234. Name: "ungrouped",
  235. Platform: PlatformAnthropic,
  236. Type: AccountTypeAPIKey,
  237. Credentials: map[string]any{"api_key": "secret"},
  238. GroupIDs: nil,
  239. }
  240. require.NoError(t, repo.Create(ctx, source))
  241. duplicate, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
  242. require.NoError(t, err)
  243. require.Empty(t, duplicate.GroupIDs)
  244. require.NotContains(t, repo.groupsOf, duplicate.ID)
  245. }
  246. func TestDuplicateAccountSimpleModeRejectsCompositeGroupBinding(t *testing.T) {
  247. ctx := context.Background()
  248. repo := newDuplicateAccountRepoStub()
  249. groupRepo := &groupRepoStubForAdmin{getByIDByID: map[int64]*Group{
  250. 9: {ID: 9, Platform: PlatformComposite},
  251. }}
  252. svc := &adminServiceImpl{
  253. cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: groupRepo,
  254. accountRepo: repo, accountDuplicateRepo: repo,
  255. }
  256. source := &Account{
  257. Name: "composite-bound", Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
  258. Credentials: map[string]any{"api_key": "secret"}, GroupIDs: []int64{9},
  259. }
  260. require.NoError(t, repo.Create(ctx, source))
  261. _, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
  262. require.Equal(t, "SIMPLE_MODE_GROUP_NOT_BINDABLE", infraerrors.Reason(err))
  263. require.Len(t, repo.accounts, 1)
  264. }
  265. func TestDuplicateAccountAtomicCreateFailureLeavesNoOrphan(t *testing.T) {
  266. ctx := context.Background()
  267. repo := newDuplicateAccountRepoStub()
  268. svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
  269. source := &Account{
  270. Name: "source",
  271. Platform: PlatformAnthropic,
  272. Type: AccountTypeAPIKey,
  273. Credentials: map[string]any{"api_key": "secret"},
  274. GroupIDs: []int64{7},
  275. AccountGroups: []AccountGroup{{GroupID: 7, Priority: 25}},
  276. }
  277. require.NoError(t, repo.Create(ctx, source))
  278. repo.atomicCreateErr = errors.New("group binding failed")
  279. _, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
  280. require.ErrorContains(t, err, "group binding failed")
  281. require.Len(t, repo.accounts, 1)
  282. }
  283. func TestDuplicateAccountNamePreservesSuffixWithinSchemaLimit(t *testing.T) {
  284. name := duplicateAccountName(strings.Repeat("界", 100))
  285. require.Equal(t, 100, utf8.RuneCountInString(name))
  286. require.True(t, strings.HasSuffix(name, " (Copy)"))
  287. }
  288. func TestDuplicateAccountReturnsExistingCopyForSameOperationKey(t *testing.T) {
  289. ctx := context.Background()
  290. repo := newDuplicateAccountRepoStub()
  291. svc := &adminServiceImpl{accountRepo: repo, accountDuplicateRepo: repo}
  292. source := &Account{
  293. Name: "source",
  294. Platform: PlatformAnthropic,
  295. Type: AccountTypeAPIKey,
  296. Credentials: map[string]any{"api_key": "secret"},
  297. }
  298. require.NoError(t, repo.Create(ctx, source))
  299. first, err := svc.DuplicateAccount(ctx, source.ID, "admin:7", "stable-operation-key")
  300. require.NoError(t, err)
  301. second, err := svc.DuplicateAccount(ctx, source.ID, "admin:7", "stable-operation-key")
  302. require.NoError(t, err)
  303. recovered, err := svc.RecoverDuplicateAccount(ctx, source.ID, "admin:7", "stable-operation-key")
  304. require.NoError(t, err)
  305. otherAdminRecovery, err := svc.RecoverDuplicateAccount(ctx, source.ID, "admin:8", "stable-operation-key")
  306. require.NoError(t, err)
  307. otherAdminCopy, err := svc.DuplicateAccount(ctx, source.ID, "admin:8", "stable-operation-key")
  308. require.NoError(t, err)
  309. require.Equal(t, first.ID, second.ID)
  310. require.Equal(t, first.ID, recovered.ID)
  311. require.Nil(t, otherAdminRecovery, "durable recovery identity must remain scoped to the initiating admin")
  312. require.NotEqual(t, first.ID, otherAdminCopy.ID)
  313. require.Len(t, repo.accounts, 3)
  314. require.NotEmpty(t, first.Extra[duplicateAccountOperationIDExtraKey])
  315. }