auth_service_platform_quota_test.go 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "fmt"
  6. "testing"
  7. "time"
  8. dbent "github.com/Wei-Shaw/sub2api/ent"
  9. )
  10. // fakeInsertRecorder 记录 BulkInsertInitial 调用,实现 UserPlatformQuotaRepository port。
  11. type fakeInsertRecorder struct {
  12. records []UserPlatformQuotaRecord
  13. err error
  14. lastCtx context.Context // 捕获最后一次 BulkInsertInitial 收到的 ctx(用于断言事务隔离)
  15. }
  16. func (f *fakeInsertRecorder) GetByUserPlatform(_ context.Context, _ int64, _ string) (*UserPlatformQuotaRecord, error) {
  17. return nil, nil
  18. }
  19. func (f *fakeInsertRecorder) BulkInsertInitial(ctx context.Context, recs []UserPlatformQuotaRecord) error {
  20. f.lastCtx = ctx
  21. if f.err != nil {
  22. return f.err
  23. }
  24. f.records = append(f.records, recs...)
  25. return nil
  26. }
  27. func (f *fakeInsertRecorder) IncrementUsageWithReset(_ context.Context, _ int64, _ string, _ float64, _ time.Time) error {
  28. return nil
  29. }
  30. func (f *fakeInsertRecorder) ListByUser(_ context.Context, _ int64) ([]UserPlatformQuotaRecord, error) {
  31. return nil, nil
  32. }
  33. func (f *fakeInsertRecorder) UpsertForUser(_ context.Context, _ int64, _ []UserPlatformQuotaRecord) error {
  34. return nil
  35. }
  36. func (f *fakeInsertRecorder) ResetExpiredWindow(_ context.Context, _ int64, _ string, _ string, _ time.Time) error {
  37. return nil
  38. }
  39. func (f *fakeInsertRecorder) BatchSnapshotUsage(_ context.Context, _ []UserPlatformQuotaSnapshot, _ time.Time) error {
  40. return nil
  41. }
  42. func TestSnapshotPlatformQuotaDefaults_PassesToRepoBulkInsert(t *testing.T) {
  43. fakeRepo := &fakeInsertRecorder{}
  44. s := &AuthService{userPlatformQuotaRepo: fakeRepo}
  45. five := 5.0
  46. plan := &signupGrantPlan{
  47. PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{
  48. "anthropic": {DailyLimitUSD: &five},
  49. "openai": {},
  50. "gemini": {},
  51. "antigravity": {},
  52. },
  53. }
  54. if err := s.snapshotPlatformQuotaDefaults(context.Background(), 999, plan); err != nil {
  55. t.Fatal(err)
  56. }
  57. if len(fakeRepo.records) != 4 {
  58. t.Errorf("expected 4 records, got %d", len(fakeRepo.records))
  59. }
  60. found := false
  61. for _, r := range fakeRepo.records {
  62. if r.UserID == 999 && r.Platform == "anthropic" && r.DailyLimitUSD != nil && *r.DailyLimitUSD == 5 {
  63. found = true
  64. }
  65. }
  66. if !found {
  67. t.Error("anthropic daily = 5 not snapshotted")
  68. }
  69. }
  70. // TestSnapshotPlatformQuotaDefaults_DetachesCallerTransaction 锁定 fix① 不变量:
  71. // 平台配额快照是 best-effort,必须脱离调用方事务执行——这样它失败(例如某平台
  72. // 违反 user_platform_quotas 的 CHECK 约束)也不会把调用方的注册主事务标记为 aborted。
  73. // 历史 bug:snapshot 在 OAuth pending handler 的 binding tx 中执行,grok 违约毒化整个
  74. // 事务 → consumePendingOAuthBrowserSessionTx 撞 "transaction aborted" → 500 → 清 cookie → 404。
  75. func TestSnapshotPlatformQuotaDefaults_DetachesCallerTransaction(t *testing.T) {
  76. fakeRepo := &fakeInsertRecorder{}
  77. s := &AuthService{userPlatformQuotaRepo: fakeRepo}
  78. five := 5.0
  79. plan := &signupGrantPlan{
  80. PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{
  81. "anthropic": {DailyLimitUSD: &five},
  82. },
  83. }
  84. // 模拟调用方(OAuth pending handler)在事务 ctx 中调用快照
  85. txCtx := dbent.NewTxContext(context.Background(), &dbent.Tx{})
  86. if err := s.snapshotPlatformQuotaDefaults(txCtx, 999, plan); err != nil {
  87. t.Fatalf("snapshot should not error (fail-open): %v", err)
  88. }
  89. if fakeRepo.lastCtx == nil {
  90. t.Fatal("expected BulkInsertInitial to be called")
  91. }
  92. if dbent.TxFromContext(fakeRepo.lastCtx) != nil {
  93. t.Error("快照必须脱离调用方事务执行(best-effort,失败不得毒化注册事务),但 repo 收到了仍携带事务的 ctx")
  94. }
  95. }
  96. func TestSnapshotPlatformQuotaDefaults_NilPlanIsNoop(t *testing.T) {
  97. fakeRepo := &fakeInsertRecorder{}
  98. s := &AuthService{userPlatformQuotaRepo: fakeRepo}
  99. if err := s.snapshotPlatformQuotaDefaults(context.Background(), 1, nil); err != nil {
  100. t.Errorf("nil plan should be noop, got %v", err)
  101. }
  102. if len(fakeRepo.records) != 0 {
  103. t.Errorf("expected no records, got %d", len(fakeRepo.records))
  104. }
  105. }
  106. func TestSnapshotPlatformQuotaDefaults_RepoErrorFailsOpen(t *testing.T) {
  107. fakeRepo := &fakeInsertRecorder{err: fmt.Errorf("db down")}
  108. s := &AuthService{userPlatformQuotaRepo: fakeRepo}
  109. five := 5.0
  110. plan := &signupGrantPlan{
  111. PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{
  112. "anthropic": {DailyLimitUSD: &five},
  113. },
  114. }
  115. if err := s.snapshotPlatformQuotaDefaults(context.Background(), 1, plan); err != nil {
  116. t.Errorf("fail-open: expected nil even on repo error, got %v", err)
  117. }
  118. }
  119. func TestSnapshotPlatformQuotaDefaults_NilRepoIsNoop(t *testing.T) {
  120. s := &AuthService{userPlatformQuotaRepo: nil}
  121. five := 5.0
  122. plan := &signupGrantPlan{
  123. PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{"a": {DailyLimitUSD: &five}},
  124. }
  125. if err := s.snapshotPlatformQuotaDefaults(context.Background(), 1, plan); err != nil {
  126. t.Errorf("nil repo should be noop, got %v", err)
  127. }
  128. }
  129. // resolveSignupGrantPlan 测试:依赖完整的 AuthService 构造,需要 SettingService(含 settingRepoStub)。
  130. // settingRepoStub 已在 auth_service_register_test.go 中定义,同 package 可直接使用。
  131. func TestResolveSignupGrantPlan_GlobalQuotaLoadedBeforeAuthSource(t *testing.T) {
  132. // 全局 quota JSON key(新格式)
  133. settings := map[string]string{
  134. SettingKeyRegistrationEnabled: "true",
  135. SettingKeyDefaultPlatformQuotas: `{
  136. "anthropic": {"daily": 10, "weekly": 50, "monthly": 200},
  137. "openai": {"daily": 5, "weekly": 25, "monthly": 100},
  138. "gemini": {"daily": 5, "weekly": 25, "monthly": 100},
  139. "antigravity": {"daily": 5, "weekly": 25, "monthly": 100}
  140. }`,
  141. }
  142. svc := newAuthService(nil, settings, nil, nil)
  143. plan := svc.resolveSignupGrantPlan(context.Background(), "email")
  144. if plan.PlatformQuotas == nil {
  145. t.Fatal("expected PlatformQuotas to be non-nil after loading global quota KVs")
  146. }
  147. q := plan.PlatformQuotas["anthropic"]
  148. if q == nil {
  149. t.Fatal("expected anthropic quota to be set")
  150. }
  151. if q.DailyLimitUSD == nil || *q.DailyLimitUSD != 10 {
  152. t.Errorf("expected anthropic daily=10, got %v", q.DailyLimitUSD)
  153. }
  154. }
  155. // TestResolveSignupGrantPlan_DisabledAuthSourceStillCarriesGlobalQuota 验证 P1 约束:
  156. // !enabled 早退路径仍携带全局 quota(GetDefaultPlatformQuotas 在 ResolveAuthSourceGrantSettings 之前)。
  157. func TestResolveSignupGrantPlan_DisabledAuthSourceStillCarriesGlobalQuota(t *testing.T) {
  158. settings := map[string]string{
  159. SettingKeyRegistrationEnabled: "true",
  160. // auth source 不配置(=> !enabled 路径)
  161. SettingKeyDefaultPlatformQuotas: `{"anthropic": {"daily": 10, "weekly": 50, "monthly": 200}}`,
  162. }
  163. svc := newAuthService(nil, settings, nil, nil)
  164. plan := svc.resolveSignupGrantPlan(context.Background(), "email")
  165. // !enabled 路径:plan.PlatformQuotas 应已含全局层(不是 nil)
  166. if plan.PlatformQuotas == nil {
  167. t.Fatal("P1 violated: PlatformQuotas is nil even with global quota KVs set")
  168. }
  169. // P1 核心断言:disabled auth source 路径不能丢失全局 quota
  170. if _, ok := plan.PlatformQuotas["anthropic"]; !ok {
  171. t.Error("P1 violated: disabled auth source path dropped global platform quota")
  172. }
  173. }