billing_cache_service_test.go 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "sync/atomic"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/stretchr/testify/require"
  10. )
  11. type billingCacheWorkerStub struct {
  12. balanceUpdates int64
  13. subscriptionUpdates int64
  14. }
  15. func (b *billingCacheWorkerStub) GetUserBalance(ctx context.Context, userID int64) (float64, error) {
  16. return 0, errors.New("not implemented")
  17. }
  18. func (b *billingCacheWorkerStub) SetUserBalance(ctx context.Context, userID int64, balance float64) error {
  19. atomic.AddInt64(&b.balanceUpdates, 1)
  20. return nil
  21. }
  22. func (b *billingCacheWorkerStub) DeductUserBalance(ctx context.Context, userID int64, amount float64) error {
  23. atomic.AddInt64(&b.balanceUpdates, 1)
  24. return nil
  25. }
  26. func (b *billingCacheWorkerStub) InvalidateUserBalance(ctx context.Context, userID int64) error {
  27. return nil
  28. }
  29. func (b *billingCacheWorkerStub) GetSubscriptionCache(ctx context.Context, userID, groupID int64) (*SubscriptionCacheData, error) {
  30. return nil, errors.New("not implemented")
  31. }
  32. func (b *billingCacheWorkerStub) SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *SubscriptionCacheData) error {
  33. atomic.AddInt64(&b.subscriptionUpdates, 1)
  34. return nil
  35. }
  36. func (b *billingCacheWorkerStub) UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error {
  37. atomic.AddInt64(&b.subscriptionUpdates, 1)
  38. return nil
  39. }
  40. func (b *billingCacheWorkerStub) InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error {
  41. return nil
  42. }
  43. func (b *billingCacheWorkerStub) GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*APIKeyRateLimitCacheData, error) {
  44. return nil, errors.New("not implemented")
  45. }
  46. func (b *billingCacheWorkerStub) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *APIKeyRateLimitCacheData) error {
  47. return nil
  48. }
  49. func (b *billingCacheWorkerStub) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
  50. return nil
  51. }
  52. func (b *billingCacheWorkerStub) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
  53. return nil
  54. }
  55. func (b *billingCacheWorkerStub) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaCacheEntry, bool, error) {
  56. return nil, false, nil
  57. }
  58. func (b *billingCacheWorkerStub) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error {
  59. return nil
  60. }
  61. func (b *billingCacheWorkerStub) DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error {
  62. return nil
  63. }
  64. func (b *billingCacheWorkerStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error {
  65. return nil
  66. }
  67. func (b *billingCacheWorkerStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) {
  68. return nil, nil
  69. }
  70. func (b *billingCacheWorkerStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error {
  71. return nil
  72. }
  73. func (b *billingCacheWorkerStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) {
  74. return nil, nil
  75. }
  76. func TestBillingCacheServiceQueueHighLoad(t *testing.T) {
  77. cache := &billingCacheWorkerStub{}
  78. svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil)
  79. t.Cleanup(svc.Stop)
  80. start := time.Now()
  81. for i := 0; i < cacheWriteBufferSize*2; i++ {
  82. svc.QueueDeductBalance(1, 1)
  83. }
  84. require.Less(t, time.Since(start), 2*time.Second)
  85. svc.QueueUpdateSubscriptionUsage(1, 2, 1.5)
  86. require.Eventually(t, func() bool {
  87. return atomic.LoadInt64(&cache.balanceUpdates) > 0
  88. }, 2*time.Second, 10*time.Millisecond)
  89. require.Eventually(t, func() bool {
  90. return atomic.LoadInt64(&cache.subscriptionUpdates) > 0
  91. }, 2*time.Second, 10*time.Millisecond)
  92. }
  93. func TestBillingCacheServiceEnqueueAfterStopReturnsFalse(t *testing.T) {
  94. cache := &billingCacheWorkerStub{}
  95. svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil)
  96. svc.Stop()
  97. enqueued := svc.enqueueCacheWrite(cacheWriteTask{
  98. kind: cacheWriteDeductBalance,
  99. userID: 1,
  100. amount: 1,
  101. })
  102. require.False(t, enqueued)
  103. }