admin_service_reset_quota_test.go 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "testing"
  7. "github.com/stretchr/testify/require"
  8. )
  9. type resetAccountQuotaRepoStub struct {
  10. mockAccountRepoForGemini
  11. account *Account
  12. getByIDErr error
  13. resetErr error
  14. resetCalls int
  15. clearRateLimitCalls int
  16. callOrder []string
  17. overloaded bool
  18. }
  19. func (r *resetAccountQuotaRepoStub) GetByID(context.Context, int64) (*Account, error) {
  20. return r.account, r.getByIDErr
  21. }
  22. func (r *resetAccountQuotaRepoStub) ResetQuotaUsedAndClearRateLimitCooldown(context.Context, int64) error {
  23. r.resetCalls++
  24. r.callOrder = append(r.callOrder, "reset_quota_and_clear_rate_limit_cooldown")
  25. return r.resetErr
  26. }
  27. func (r *resetAccountQuotaRepoStub) ClearRateLimit(context.Context, int64) error {
  28. r.clearRateLimitCalls++
  29. r.callOrder = append(r.callOrder, "clear_rate_limit")
  30. r.overloaded = false
  31. return nil
  32. }
  33. func TestResetAccountQuota_ClearsSchedulerRateLimitWithoutClearingOverload(t *testing.T) {
  34. repo := &resetAccountQuotaRepoStub{account: &Account{ID: 42}, overloaded: true}
  35. svc := &adminServiceImpl{accountRepo: repo}
  36. err := svc.ResetAccountQuota(context.Background(), 42)
  37. require.NoError(t, err)
  38. require.Equal(t, 1, repo.resetCalls)
  39. require.Zero(t, repo.clearRateLimitCalls)
  40. require.Equal(t, []string{"reset_quota_and_clear_rate_limit_cooldown"}, repo.callOrder)
  41. require.True(t, repo.overloaded, "quota reset must preserve an unrelated overload block")
  42. }
  43. func TestResetAccountQuota_PreservesLookupAndSparkShadowShortCircuits(t *testing.T) {
  44. t.Run("lookup failure", func(t *testing.T) {
  45. getErr := errors.New("get account failed")
  46. repo := &resetAccountQuotaRepoStub{getByIDErr: getErr}
  47. svc := &adminServiceImpl{accountRepo: repo}
  48. err := svc.ResetAccountQuota(context.Background(), 42)
  49. require.ErrorIs(t, err, getErr)
  50. require.Zero(t, repo.resetCalls)
  51. require.Zero(t, repo.clearRateLimitCalls)
  52. })
  53. t.Run("spark shadow", func(t *testing.T) {
  54. parentID := int64(7)
  55. repo := &resetAccountQuotaRepoStub{
  56. account: &Account{ID: 42, ParentAccountID: &parentID},
  57. }
  58. svc := &adminServiceImpl{accountRepo: repo}
  59. err := svc.ResetAccountQuota(context.Background(), 42)
  60. require.Error(t, err)
  61. require.Zero(t, repo.resetCalls)
  62. require.Zero(t, repo.clearRateLimitCalls)
  63. })
  64. }
  65. func TestResetAccountQuota_PropagatesAtomicRepositoryFailure(t *testing.T) {
  66. resetErr := errors.New("atomic reset failed")
  67. repo := &resetAccountQuotaRepoStub{
  68. account: &Account{ID: 42},
  69. resetErr: resetErr,
  70. }
  71. svc := &adminServiceImpl{accountRepo: repo}
  72. err := svc.ResetAccountQuota(context.Background(), 42)
  73. require.ErrorIs(t, err, resetErr)
  74. require.Equal(t, 1, repo.resetCalls)
  75. require.Zero(t, repo.clearRateLimitCalls)
  76. }