billing_cache_service_rpm_test.go 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "sync/atomic"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/stretchr/testify/require"
  10. )
  11. // userRPMCacheStub 记录每种计数器被调用的次数,并可注入返回值与错误。
  12. type userRPMCacheStub struct {
  13. userGroupCalls int32
  14. userCalls int32
  15. userGroupCounts []int // 依次返回的计数值
  16. userGroupErr error
  17. userCounts []int
  18. userErr error
  19. }
  20. func (s *userRPMCacheStub) IncrementUserGroupRPM(_ context.Context, _, _ int64) (int, error) {
  21. idx := int(atomic.AddInt32(&s.userGroupCalls, 1)) - 1
  22. if s.userGroupErr != nil {
  23. return 0, s.userGroupErr
  24. }
  25. if idx < len(s.userGroupCounts) {
  26. return s.userGroupCounts[idx], nil
  27. }
  28. return 1, nil
  29. }
  30. func (s *userRPMCacheStub) IncrementUserRPM(_ context.Context, _ int64) (int, error) {
  31. idx := int(atomic.AddInt32(&s.userCalls, 1)) - 1
  32. if s.userErr != nil {
  33. return 0, s.userErr
  34. }
  35. if idx < len(s.userCounts) {
  36. return s.userCounts[idx], nil
  37. }
  38. return 1, nil
  39. }
  40. func (s *userRPMCacheStub) GetUserGroupRPM(_ context.Context, _, _ int64) (int, error) {
  41. return 0, nil
  42. }
  43. func (s *userRPMCacheStub) GetUserRPM(_ context.Context, _ int64) (int, error) {
  44. return 0, nil
  45. }
  46. // rpmOverrideRepoStub 专用于 checkRPM 分支测试,只实现必要方法。
  47. type rpmOverrideRepoStub struct {
  48. UserGroupRateRepository
  49. override *int
  50. err error
  51. calls int32
  52. }
  53. func (s *rpmOverrideRepoStub) GetRPMOverrideByUserAndGroup(_ context.Context, _, _ int64) (*int, error) {
  54. atomic.AddInt32(&s.calls, 1)
  55. if s.err != nil {
  56. return nil, s.err
  57. }
  58. return s.override, nil
  59. }
  60. func newBillingServiceForRPM(t *testing.T, cache UserRPMCache, rateRepo UserGroupRateRepository) *BillingCacheService {
  61. t.Helper()
  62. // 用 nil BillingCache 走 "无缓存" 分支,避免 CheckBillingEligibility 副作用。
  63. // 我们只直接测 checkRPM。
  64. svc := NewBillingCacheService(nil, nil, nil, nil, cache, rateRepo, &config.Config{}, nil)
  65. t.Cleanup(svc.Stop)
  66. return svc
  67. }
  68. func TestBillingCacheService_CheckRPM_OverrideTakesPrecedenceOverGroup(t *testing.T) {
  69. override := 2
  70. // user-group 计数: 1, 2, 3;user 计数: 默认返回 1(远小于 RPMLimit=100,不干扰)
  71. cache := &userRPMCacheStub{userGroupCounts: []int{1, 2, 3}}
  72. repo := &rpmOverrideRepoStub{override: &override}
  73. svc := newBillingServiceForRPM(t, cache, repo)
  74. user := &User{ID: 1, RPMLimit: 100} // 全局上限设高,不干扰 override 测试
  75. group := &Group{ID: 10, RPMLimit: 100}
  76. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  77. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  78. require.ErrorIs(t, svc.checkRPM(context.Background(), user, group), ErrGroupRPMExceeded)
  79. require.EqualValues(t, 3, atomic.LoadInt32(&cache.userGroupCalls), "override 命中分支应走 user-group 计数")
  80. // 并行设计:前 2 次 override 未超→继续检查 user;第 3 次 override 超了→直接 return,不检查 user
  81. require.EqualValues(t, 2, atomic.LoadInt32(&cache.userCalls), "override 超限前 user 计数器应被调用")
  82. require.EqualValues(t, 3, atomic.LoadInt32(&repo.calls))
  83. }
  84. func TestBillingCacheService_CheckRPM_UserLimitIsGlobalHardCap(t *testing.T) {
  85. override := 100 // override 很高
  86. // user-group 计数: 默认返回 1(远小于 override);user 计数: 1, 2, 3
  87. cache := &userRPMCacheStub{userCounts: []int{1, 2, 3}}
  88. repo := &rpmOverrideRepoStub{override: &override}
  89. svc := newBillingServiceForRPM(t, cache, repo)
  90. user := &User{ID: 1, RPMLimit: 2} // 全局硬上限=2,应覆盖 override=100
  91. group := &Group{ID: 10, RPMLimit: 100}
  92. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  93. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  94. require.ErrorIs(t, svc.checkRPM(context.Background(), user, group), ErrUserRPMExceeded, "user 全局硬上限应优先于 override")
  95. }
  96. func TestBillingCacheService_CheckRPM_OverrideZeroSkipsGroupButUserStillApplies(t *testing.T) {
  97. zero := 0
  98. // user 计数: 依次返回 1..6
  99. cache := &userRPMCacheStub{userCounts: []int{1, 2, 3, 4, 5, 6}}
  100. repo := &rpmOverrideRepoStub{override: &zero}
  101. svc := newBillingServiceForRPM(t, cache, repo)
  102. user := &User{ID: 1, RPMLimit: 5}
  103. group := &Group{ID: 10, RPMLimit: 100}
  104. // override=0 跳过分组计数,但 user.RPMLimit=5 仍生效
  105. for i := 0; i < 5; i++ {
  106. require.NoError(t, svc.checkRPM(context.Background(), user, group), "request %d should pass", i+1)
  107. }
  108. require.ErrorIs(t, svc.checkRPM(context.Background(), user, group), ErrUserRPMExceeded,
  109. "override=0 跳过分组但 user 全局上限仍应生效")
  110. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userGroupCalls), "override=0 不应触发分组计数器")
  111. require.EqualValues(t, 6, atomic.LoadInt32(&cache.userCalls), "user 计数器应被调用")
  112. }
  113. func TestBillingCacheService_CheckRPM_OverrideZeroAndUserZeroIsFullyUnlimited(t *testing.T) {
  114. zero := 0
  115. cache := &userRPMCacheStub{}
  116. repo := &rpmOverrideRepoStub{override: &zero}
  117. svc := newBillingServiceForRPM(t, cache, repo)
  118. user := &User{ID: 1, RPMLimit: 0} // user 也不限
  119. group := &Group{ID: 10, RPMLimit: 100}
  120. for i := 0; i < 50; i++ {
  121. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  122. }
  123. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userGroupCalls), "override=0 不触发分组计数")
  124. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userCalls), "user.RPMLimit=0 也不触发用户计数")
  125. }
  126. func TestBillingCacheService_CheckRPM_NilOverrideFallsThroughToGroup(t *testing.T) {
  127. // user-group 计数: 5, 6;user 计数: 默认 1(不干扰)
  128. cache := &userRPMCacheStub{userGroupCounts: []int{5, 6}}
  129. repo := &rpmOverrideRepoStub{override: nil}
  130. svc := newBillingServiceForRPM(t, cache, repo)
  131. user := &User{ID: 1, RPMLimit: 999} // 全局上限很高,group 先超
  132. group := &Group{ID: 10, RPMLimit: 5}
  133. require.NoError(t, svc.checkRPM(context.Background(), user, group)) // ug=5, user=1, 都没超
  134. require.ErrorIs(t, svc.checkRPM(context.Background(), user, group), ErrGroupRPMExceeded) // ug=6 > 5
  135. require.EqualValues(t, 2, atomic.LoadInt32(&cache.userGroupCalls))
  136. // 并行模式:第 1 次 group 没超 → 继续检查 user;第 2 次 group 超了 → 直接 return,不检查 user
  137. require.EqualValues(t, 1, atomic.LoadInt32(&cache.userCalls), "group 未超时 user 也应检查;group 超时直接返回")
  138. }
  139. func TestBillingCacheService_CheckRPM_OverrideLookupErrorFallsThroughToGroup(t *testing.T) {
  140. cache := &userRPMCacheStub{userGroupCounts: []int{3}}
  141. repo := &rpmOverrideRepoStub{err: errors.New("db down")}
  142. svc := newBillingServiceForRPM(t, cache, repo)
  143. user := &User{ID: 1, RPMLimit: 0}
  144. group := &Group{ID: 10, RPMLimit: 10}
  145. // override 查询失败后应继续尝试 group 分支(不直接拒绝)
  146. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  147. require.EqualValues(t, 1, atomic.LoadInt32(&cache.userGroupCalls))
  148. require.EqualValues(t, 1, atomic.LoadInt32(&repo.calls))
  149. }
  150. func TestBillingCacheService_CheckRPM_UserLevelFallbackWhenGroupUnlimited(t *testing.T) {
  151. cache := &userRPMCacheStub{userCounts: []int{1, 2, 3}}
  152. repo := &rpmOverrideRepoStub{override: nil}
  153. svc := newBillingServiceForRPM(t, cache, repo)
  154. user := &User{ID: 1, RPMLimit: 2}
  155. group := &Group{ID: 10, RPMLimit: 0} // 分组未设限
  156. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  157. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  158. require.ErrorIs(t, svc.checkRPM(context.Background(), user, group), ErrUserRPMExceeded)
  159. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userGroupCalls), "group 未设限时不应 INCR user-group 键")
  160. require.EqualValues(t, 3, atomic.LoadInt32(&cache.userCalls))
  161. }
  162. func TestBillingCacheService_CheckRPM_NoLimitsConfiguredIsNoop(t *testing.T) {
  163. cache := &userRPMCacheStub{}
  164. repo := &rpmOverrideRepoStub{override: nil}
  165. svc := newBillingServiceForRPM(t, cache, repo)
  166. user := &User{ID: 1, RPMLimit: 0}
  167. group := &Group{ID: 10, RPMLimit: 0}
  168. for i := 0; i < 10; i++ {
  169. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  170. }
  171. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userGroupCalls))
  172. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userCalls))
  173. }
  174. func TestBillingCacheService_CheckRPM_RedisErrorFailOpen(t *testing.T) {
  175. cache := &userRPMCacheStub{userGroupErr: errors.New("redis unavailable")}
  176. repo := &rpmOverrideRepoStub{override: nil}
  177. svc := newBillingServiceForRPM(t, cache, repo)
  178. user := &User{ID: 1, RPMLimit: 0}
  179. group := &Group{ID: 10, RPMLimit: 5}
  180. // Redis 故障时应 fail-open,不拒绝请求
  181. require.NoError(t, svc.checkRPM(context.Background(), user, group))
  182. require.EqualValues(t, 1, atomic.LoadInt32(&cache.userGroupCalls))
  183. }
  184. func TestBillingCacheService_CheckRPM_NoGroupUsesUserOnly(t *testing.T) {
  185. cache := &userRPMCacheStub{userCounts: []int{1, 2, 3}}
  186. repo := &rpmOverrideRepoStub{}
  187. svc := newBillingServiceForRPM(t, cache, repo)
  188. user := &User{ID: 1, RPMLimit: 2}
  189. // 无 group(纯用户级限流场景),不应查询 rpm_override。
  190. require.NoError(t, svc.checkRPM(context.Background(), user, nil))
  191. require.NoError(t, svc.checkRPM(context.Background(), user, nil))
  192. require.ErrorIs(t, svc.checkRPM(context.Background(), user, nil), ErrUserRPMExceeded)
  193. require.EqualValues(t, 0, atomic.LoadInt32(&repo.calls), "无 group 时不应查询 rpm_override")
  194. require.EqualValues(t, 3, atomic.LoadInt32(&cache.userCalls))
  195. }
  196. func TestBillingCacheService_CheckRPM_NilUserIsNoop(t *testing.T) {
  197. cache := &userRPMCacheStub{}
  198. repo := &rpmOverrideRepoStub{}
  199. svc := newBillingServiceForRPM(t, cache, repo)
  200. require.NoError(t, svc.checkRPM(context.Background(), nil, &Group{ID: 1, RPMLimit: 10}))
  201. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userGroupCalls))
  202. require.EqualValues(t, 0, atomic.LoadInt32(&cache.userCalls))
  203. require.EqualValues(t, 0, atomic.LoadInt32(&repo.calls))
  204. }