ratelimit_service_openai_team_linked_test.go 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "net/http"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/stretchr/testify/require"
  10. )
  11. const teamLinkedDeactivatedBody = `{"detail":{"code":"deactivated_workspace","message":"This workspace has been deactivated."}}`
  12. type teamLinkedAccountRepoStub struct {
  13. mockAccountRepoForGemini
  14. teamAccounts []Account
  15. listErr error
  16. listCalls int
  17. setErrorIDs []int64
  18. setErrorMsgs map[int64]string
  19. failSetError map[int64]error
  20. }
  21. // ListByPlatform 镜像真实仓库语义:仅返回该平台的 active 账户。
  22. func (r *teamLinkedAccountRepoStub) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
  23. r.listCalls++
  24. if r.listErr != nil {
  25. return nil, r.listErr
  26. }
  27. out := make([]Account, 0, len(r.teamAccounts))
  28. for _, acc := range r.teamAccounts {
  29. if acc.Platform == platform && acc.Status == StatusActive {
  30. out = append(out, acc)
  31. }
  32. }
  33. return out, nil
  34. }
  35. func (r *teamLinkedAccountRepoStub) SetError(ctx context.Context, id int64, errorMsg string) error {
  36. if err, ok := r.failSetError[id]; ok {
  37. return err
  38. }
  39. r.setErrorIDs = append(r.setErrorIDs, id)
  40. if r.setErrorMsgs == nil {
  41. r.setErrorMsgs = make(map[int64]string)
  42. }
  43. r.setErrorMsgs[id] = errorMsg
  44. return nil
  45. }
  46. func newTeamLinkedAccount(id int64, teamID string) Account {
  47. return Account{
  48. ID: id,
  49. Platform: PlatformOpenAI,
  50. Type: AccountTypeOAuth,
  51. Status: StatusActive,
  52. Credentials: map[string]any{"chatgpt_account_id": teamID},
  53. }
  54. }
  55. // newTeamLinkedFixture: #1 触发者(team-A) #2 同队 #3 异队 #4 apikey #5 影子 #6 同队 #7 同队但已 error
  56. func newTeamLinkedFixture() []Account {
  57. parentID := int64(1)
  58. shadow := Account{
  59. ID: 5,
  60. Platform: PlatformOpenAI,
  61. Type: AccountTypeOAuth,
  62. Status: StatusActive,
  63. ParentAccountID: &parentID,
  64. }
  65. apikey := newTeamLinkedAccount(4, "team-A")
  66. apikey.Type = AccountTypeAPIKey
  67. erroredSibling := newTeamLinkedAccount(7, "team-A")
  68. erroredSibling.Status = StatusError
  69. return []Account{
  70. newTeamLinkedAccount(1, "team-A"),
  71. newTeamLinkedAccount(2, "team-A"),
  72. newTeamLinkedAccount(3, "team-B"),
  73. apikey,
  74. shadow,
  75. newTeamLinkedAccount(6, "team-A"),
  76. erroredSibling,
  77. }
  78. }
  79. func newTeamLinkedTestService(repo *teamLinkedAccountRepoStub) (*RateLimitService, *runtimeBlockRecorder) {
  80. rl := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
  81. blocker := &runtimeBlockRecorder{}
  82. rl.SetAccountRuntimeBlocker(blocker)
  83. return rl, blocker
  84. }
  85. func TestTeamLinkedError_FanoutMarksSameTeamAccounts(t *testing.T) {
  86. repo := &teamLinkedAccountRepoStub{teamAccounts: newTeamLinkedFixture()}
  87. rl, blocker := newTeamLinkedTestService(repo)
  88. trigger := newTeamLinkedAccount(1, "team-A")
  89. shouldDisable := rl.HandleUpstreamError(context.Background(), &trigger, http.StatusPaymentRequired, http.Header{}, []byte(teamLinkedDeactivatedBody))
  90. require.True(t, shouldDisable)
  91. // fan-out 先标记同队兄弟(#2、#6),触发账户 #1 随后由常规 case 402 标记
  92. require.Equal(t, []int64{2, 6, 1}, repo.setErrorIDs)
  93. require.Contains(t, repo.setErrorMsgs[2], "team-linked error triggered by account #1")
  94. require.Contains(t, repo.setErrorMsgs[6], "team-linked error triggered by account #1")
  95. require.Contains(t, repo.setErrorMsgs[1], "Workspace deactivated (402)")
  96. require.NotContains(t, repo.setErrorMsgs[1], "team-linked")
  97. // 熔断顺序:兄弟账户先于落库全部进程内熔断,触发账户走 auth_error
  98. require.Equal(t, []string{openAITeamLinkedErrorBlockReason, openAITeamLinkedErrorBlockReason, "auth_error"}, blocker.reasons)
  99. require.Equal(t, int64(2), blocker.accounts[0].ID)
  100. require.Equal(t, int64(6), blocker.accounts[1].ID)
  101. require.Equal(t, int64(1), blocker.accounts[2].ID)
  102. }
  103. func TestTeamLinkedError_GenericPaymentErrorDoesNotFanout(t *testing.T) {
  104. repo := &teamLinkedAccountRepoStub{teamAccounts: newTeamLinkedFixture()}
  105. rl, _ := newTeamLinkedTestService(repo)
  106. trigger := newTeamLinkedAccount(1, "team-A")
  107. rl.HandleUpstreamError(context.Background(), &trigger, http.StatusPaymentRequired, http.Header{}, []byte(`{"error":{"message":"insufficient balance"}}`))
  108. require.Equal(t, []int64{1}, repo.setErrorIDs)
  109. require.Contains(t, repo.setErrorMsgs[1], "Payment required (402)")
  110. require.Zero(t, repo.listCalls)
  111. }
  112. func TestTeamLinkedError_DedupWithinTTL(t *testing.T) {
  113. repo := &teamLinkedAccountRepoStub{teamAccounts: newTeamLinkedFixture()}
  114. rl, _ := newTeamLinkedTestService(repo)
  115. first := newTeamLinkedAccount(1, "team-A")
  116. second := newTeamLinkedAccount(2, "team-A")
  117. rl.HandleUpstreamError(context.Background(), &first, http.StatusPaymentRequired, http.Header{}, []byte(teamLinkedDeactivatedBody))
  118. rl.HandleUpstreamError(context.Background(), &second, http.StatusPaymentRequired, http.Header{}, []byte(teamLinkedDeactivatedBody))
  119. // 第二次触发被去重:只有 #2 自身经 case 402 标记,未再次 fan-out
  120. require.Equal(t, []int64{2, 6, 1, 2}, repo.setErrorIDs)
  121. require.Equal(t, 1, repo.listCalls)
  122. }
  123. func TestTeamLinkedError_APIKeyTriggerDoesNotFanout(t *testing.T) {
  124. repo := &teamLinkedAccountRepoStub{teamAccounts: newTeamLinkedFixture()}
  125. rl, _ := newTeamLinkedTestService(repo)
  126. trigger := newTeamLinkedAccount(4, "team-A")
  127. trigger.Type = AccountTypeAPIKey
  128. rl.HandleUpstreamError(context.Background(), &trigger, http.StatusPaymentRequired, http.Header{}, []byte(teamLinkedDeactivatedBody))
  129. require.Equal(t, []int64{4}, repo.setErrorIDs)
  130. require.Zero(t, repo.listCalls)
  131. }
  132. func TestTeamLinkedError_DirectCallSkipsTriggerAccount(t *testing.T) {
  133. // 直调对应 fastpath 调用点:账户级临时不可调度规则短路时联动仍然生效
  134. repo := &teamLinkedAccountRepoStub{teamAccounts: newTeamLinkedFixture()}
  135. rl, blocker := newTeamLinkedTestService(repo)
  136. trigger := newTeamLinkedAccount(1, "team-A")
  137. rl.maybeHandleOpenAITeamLinkedError(context.Background(), &trigger, http.StatusPaymentRequired, []byte(teamLinkedDeactivatedBody))
  138. require.Equal(t, []int64{2, 6}, repo.setErrorIDs)
  139. require.Equal(t, []string{openAITeamLinkedErrorBlockReason, openAITeamLinkedErrorBlockReason}, blocker.reasons)
  140. }
  141. func TestTeamLinkedError_MissingTeamIDDoesNothing(t *testing.T) {
  142. repo := &teamLinkedAccountRepoStub{teamAccounts: newTeamLinkedFixture()}
  143. rl, blocker := newTeamLinkedTestService(repo)
  144. trigger := Account{ID: 9, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive}
  145. rl.maybeHandleOpenAITeamLinkedError(context.Background(), &trigger, http.StatusPaymentRequired, []byte(teamLinkedDeactivatedBody))
  146. require.Empty(t, repo.setErrorIDs)
  147. require.Empty(t, blocker.reasons)
  148. require.Zero(t, repo.listCalls)
  149. }
  150. func TestTeamLinkedError_SetErrorFailureDoesNotAbortRemaining(t *testing.T) {
  151. repo := &teamLinkedAccountRepoStub{
  152. teamAccounts: newTeamLinkedFixture(),
  153. failSetError: map[int64]error{2: errors.New("db down")},
  154. }
  155. rl, blocker := newTeamLinkedTestService(repo)
  156. trigger := newTeamLinkedAccount(1, "team-A")
  157. rl.maybeHandleOpenAITeamLinkedError(context.Background(), &trigger, http.StatusPaymentRequired, []byte(teamLinkedDeactivatedBody))
  158. require.Equal(t, []int64{6}, repo.setErrorIDs)
  159. // 进程内熔断先于落库执行,两个账户都已被熔断
  160. require.Len(t, blocker.reasons, 2)
  161. }