subscription_reset_quota_test.go 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
  9. "github.com/stretchr/testify/require"
  10. )
  11. // resetQuotaUserSubRepoStub 支持 GetByID、ResetUsageWindows,
  12. // 其余方法继承 userSubRepoNoop(panic)。
  13. type resetQuotaUserSubRepoStub struct {
  14. userSubRepoNoop
  15. sub *UserSubscription
  16. resetDailyCalled bool
  17. resetWeeklyCalled bool
  18. resetMonthlyCalled bool
  19. resetDailyErr error
  20. resetWeeklyErr error
  21. resetMonthlyErr error
  22. dailyStart time.Time
  23. periodicStart time.Time
  24. }
  25. func (r *resetQuotaUserSubRepoStub) GetByID(_ context.Context, id int64) (*UserSubscription, error) {
  26. if r.sub == nil || r.sub.ID != id {
  27. return nil, ErrSubscriptionNotFound
  28. }
  29. cp := *r.sub
  30. return &cp, nil
  31. }
  32. func (r *resetQuotaUserSubRepoStub) ResetUsageWindows(_ context.Context, _ int64, resetDaily, resetWeekly, resetMonthly bool, dailyStart, periodicStart time.Time) error {
  33. r.resetDailyCalled = resetDaily
  34. r.resetWeeklyCalled = resetWeekly
  35. r.resetMonthlyCalled = resetMonthly
  36. r.dailyStart = dailyStart
  37. r.periodicStart = periodicStart
  38. if resetDaily && r.resetDailyErr != nil {
  39. return r.resetDailyErr
  40. }
  41. if resetWeekly && r.resetWeeklyErr != nil {
  42. return r.resetWeeklyErr
  43. }
  44. if resetMonthly && r.resetMonthlyErr != nil {
  45. return r.resetMonthlyErr
  46. }
  47. if r.sub == nil {
  48. return nil
  49. }
  50. if resetDaily {
  51. r.sub.DailyUsageUSD = 0
  52. r.sub.DailyWindowStart = &dailyStart
  53. }
  54. if resetWeekly {
  55. r.sub.WeeklyUsageUSD = 0
  56. r.sub.WeeklyWindowStart = &periodicStart
  57. }
  58. if resetMonthly {
  59. r.sub.MonthlyUsageUSD = 0
  60. r.sub.MonthlyWindowStart = &periodicStart
  61. }
  62. return nil
  63. }
  64. func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, _ *time.Time, windowStart time.Time) error {
  65. r.resetDailyCalled = true
  66. if r.resetDailyErr == nil && r.sub != nil {
  67. r.sub.DailyUsageUSD = 0
  68. r.sub.DailyWindowStart = &windowStart
  69. }
  70. return r.resetDailyErr
  71. }
  72. func (r *resetQuotaUserSubRepoStub) ResetWeeklyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
  73. r.resetWeeklyCalled = true
  74. return r.resetWeeklyErr
  75. }
  76. func (r *resetQuotaUserSubRepoStub) ResetMonthlyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
  77. r.resetMonthlyCalled = true
  78. return r.resetMonthlyErr
  79. }
  80. func newResetQuotaSvc(stub *resetQuotaUserSubRepoStub) *SubscriptionService {
  81. return NewSubscriptionService(groupRepoNoop{}, stub, nil, nil, nil)
  82. }
  83. func TestAdminResetQuota_ResetBoth(t *testing.T) {
  84. stub := &resetQuotaUserSubRepoStub{
  85. sub: &UserSubscription{ID: 1, UserID: 10, GroupID: 20},
  86. }
  87. svc := newResetQuotaSvc(stub)
  88. resetAt := time.Date(2026, 7, 1, 10, 37, 42, 123, time.UTC)
  89. svc.now = func() time.Time { return resetAt }
  90. result, err := svc.AdminResetQuota(context.Background(), 1, true, true, false)
  91. require.NoError(t, err)
  92. require.NotNil(t, result)
  93. require.True(t, stub.resetDailyCalled, "应调用 ResetDailyUsage")
  94. require.True(t, stub.resetWeeklyCalled, "应调用 ResetWeeklyUsage")
  95. require.False(t, stub.resetMonthlyCalled, "不应调用 ResetMonthlyUsage")
  96. // 手动重置后日窗口锚定当天 0 点(保持 0 点刷新节奏),周窗口锚定重置时刻。
  97. require.Equal(t, timezone.StartOfDay(resetAt), stub.dailyStart)
  98. require.Equal(t, resetAt, stub.periodicStart)
  99. require.Equal(t, timezone.StartOfDay(resetAt), *result.DailyWindowStart)
  100. require.Equal(t, resetAt, *result.WeeklyWindowStart)
  101. }
  102. func TestAdminResetQuota_ResetDailyOnly(t *testing.T) {
  103. stub := &resetQuotaUserSubRepoStub{
  104. sub: &UserSubscription{ID: 2, UserID: 10, GroupID: 20},
  105. }
  106. svc := newResetQuotaSvc(stub)
  107. result, err := svc.AdminResetQuota(context.Background(), 2, true, false, false)
  108. require.NoError(t, err)
  109. require.NotNil(t, result)
  110. require.True(t, stub.resetDailyCalled, "应调用 ResetDailyUsage")
  111. require.False(t, stub.resetWeeklyCalled, "不应调用 ResetWeeklyUsage")
  112. require.False(t, stub.resetMonthlyCalled, "不应调用 ResetMonthlyUsage")
  113. }
  114. func TestAdminResetQuota_ResetWeeklyOnly(t *testing.T) {
  115. stub := &resetQuotaUserSubRepoStub{
  116. sub: &UserSubscription{ID: 3, UserID: 10, GroupID: 20},
  117. }
  118. svc := newResetQuotaSvc(stub)
  119. result, err := svc.AdminResetQuota(context.Background(), 3, false, true, false)
  120. require.NoError(t, err)
  121. require.NotNil(t, result)
  122. require.False(t, stub.resetDailyCalled, "不应调用 ResetDailyUsage")
  123. require.True(t, stub.resetWeeklyCalled, "应调用 ResetWeeklyUsage")
  124. require.False(t, stub.resetMonthlyCalled, "不应调用 ResetMonthlyUsage")
  125. }
  126. func TestAdminResetQuota_BothFalseReturnsError(t *testing.T) {
  127. stub := &resetQuotaUserSubRepoStub{
  128. sub: &UserSubscription{ID: 7, UserID: 10, GroupID: 20},
  129. }
  130. svc := newResetQuotaSvc(stub)
  131. _, err := svc.AdminResetQuota(context.Background(), 7, false, false, false)
  132. require.ErrorIs(t, err, ErrInvalidInput)
  133. require.False(t, stub.resetDailyCalled)
  134. require.False(t, stub.resetWeeklyCalled)
  135. require.False(t, stub.resetMonthlyCalled)
  136. }
  137. func TestAdminResetQuota_SubscriptionNotFound(t *testing.T) {
  138. stub := &resetQuotaUserSubRepoStub{sub: nil}
  139. svc := newResetQuotaSvc(stub)
  140. _, err := svc.AdminResetQuota(context.Background(), 999, true, true, true)
  141. require.ErrorIs(t, err, ErrSubscriptionNotFound)
  142. require.False(t, stub.resetDailyCalled)
  143. require.False(t, stub.resetWeeklyCalled)
  144. require.False(t, stub.resetMonthlyCalled)
  145. }
  146. func TestAdminResetQuota_ResetDailyUsageError(t *testing.T) {
  147. dbErr := errors.New("db error")
  148. stub := &resetQuotaUserSubRepoStub{
  149. sub: &UserSubscription{ID: 4, UserID: 10, GroupID: 20},
  150. resetDailyErr: dbErr,
  151. }
  152. svc := newResetQuotaSvc(stub)
  153. _, err := svc.AdminResetQuota(context.Background(), 4, true, true, false)
  154. require.ErrorIs(t, err, dbErr)
  155. require.True(t, stub.resetDailyCalled)
  156. require.True(t, stub.resetWeeklyCalled, "原子重置应在一次调用中提交所选窗口")
  157. }
  158. func TestAdminResetQuota_ResetWeeklyUsageError(t *testing.T) {
  159. dbErr := errors.New("db error")
  160. stub := &resetQuotaUserSubRepoStub{
  161. sub: &UserSubscription{ID: 5, UserID: 10, GroupID: 20},
  162. resetWeeklyErr: dbErr,
  163. }
  164. svc := newResetQuotaSvc(stub)
  165. _, err := svc.AdminResetQuota(context.Background(), 5, false, true, false)
  166. require.ErrorIs(t, err, dbErr)
  167. require.True(t, stub.resetWeeklyCalled)
  168. }
  169. func TestAdminResetQuota_ResetMonthlyOnly(t *testing.T) {
  170. stub := &resetQuotaUserSubRepoStub{
  171. sub: &UserSubscription{ID: 8, UserID: 10, GroupID: 20},
  172. }
  173. svc := newResetQuotaSvc(stub)
  174. result, err := svc.AdminResetQuota(context.Background(), 8, false, false, true)
  175. require.NoError(t, err)
  176. require.NotNil(t, result)
  177. require.False(t, stub.resetDailyCalled, "不应调用 ResetDailyUsage")
  178. require.False(t, stub.resetWeeklyCalled, "不应调用 ResetWeeklyUsage")
  179. require.True(t, stub.resetMonthlyCalled, "应调用 ResetMonthlyUsage")
  180. }
  181. func TestAdminResetQuota_BeforeStartsAtSameDayPreservesAutomaticBoundary(t *testing.T) {
  182. startsAt := time.Date(2026, 7, 1, 15, 0, 0, 0, time.UTC)
  183. resetAt := time.Date(2026, 7, 1, 10, 37, 42, 123, time.UTC)
  184. stub := &resetQuotaUserSubRepoStub{
  185. sub: &UserSubscription{
  186. ID: 10,
  187. UserID: 10,
  188. GroupID: 20,
  189. StartsAt: startsAt,
  190. ExpiresAt: startsAt.Add(45 * 24 * time.Hour),
  191. },
  192. }
  193. svc := newResetQuotaSvc(stub)
  194. svc.now = func() time.Time { return resetAt }
  195. result, err := svc.AdminResetQuota(context.Background(), 10, false, false, true)
  196. require.NoError(t, err)
  197. require.Equal(t, resetAt, *result.MonthlyWindowStart)
  198. boundary, ok := result.automaticWindowStartAt(result.MonthlyWindowStart, 30*24*time.Hour, resetAt.Add(30*24*time.Hour))
  199. require.True(t, ok)
  200. require.Equal(t, resetAt.Add(30*24*time.Hour), boundary)
  201. }
  202. func TestAdminResetQuota_ResetMonthlyUsageError(t *testing.T) {
  203. dbErr := errors.New("db error")
  204. stub := &resetQuotaUserSubRepoStub{
  205. sub: &UserSubscription{ID: 9, UserID: 10, GroupID: 20},
  206. resetMonthlyErr: dbErr,
  207. }
  208. svc := newResetQuotaSvc(stub)
  209. _, err := svc.AdminResetQuota(context.Background(), 9, false, false, true)
  210. require.ErrorIs(t, err, dbErr)
  211. require.True(t, stub.resetMonthlyCalled)
  212. }
  213. func TestAdminResetQuota_ReturnsRefreshedSub(t *testing.T) {
  214. stub := &resetQuotaUserSubRepoStub{
  215. sub: &UserSubscription{
  216. ID: 6,
  217. UserID: 10,
  218. GroupID: 20,
  219. DailyUsageUSD: 99.9,
  220. },
  221. }
  222. svc := newResetQuotaSvc(stub)
  223. result, err := svc.AdminResetQuota(context.Background(), 6, true, false, false)
  224. require.NoError(t, err)
  225. // ResetUsageWindows stub 会将 sub.DailyUsageUSD 归零,
  226. // 服务应返回第二次 GetByID 的刷新值而非初始的 99.9
  227. require.Equal(t, float64(0), result.DailyUsageUSD, "返回的订阅应反映已归零的用量")
  228. require.True(t, stub.resetDailyCalled)
  229. }