| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132 |
- package service
- import (
- "context"
- "errors"
- "sync/atomic"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/stretchr/testify/require"
- )
- type billingCacheWorkerStub struct {
- balanceUpdates int64
- subscriptionUpdates int64
- }
- func (b *billingCacheWorkerStub) GetUserBalance(ctx context.Context, userID int64) (float64, error) {
- return 0, errors.New("not implemented")
- }
- func (b *billingCacheWorkerStub) SetUserBalance(ctx context.Context, userID int64, balance float64) error {
- atomic.AddInt64(&b.balanceUpdates, 1)
- return nil
- }
- func (b *billingCacheWorkerStub) DeductUserBalance(ctx context.Context, userID int64, amount float64) error {
- atomic.AddInt64(&b.balanceUpdates, 1)
- return nil
- }
- func (b *billingCacheWorkerStub) InvalidateUserBalance(ctx context.Context, userID int64) error {
- return nil
- }
- func (b *billingCacheWorkerStub) GetSubscriptionCache(ctx context.Context, userID, groupID int64) (*SubscriptionCacheData, error) {
- return nil, errors.New("not implemented")
- }
- func (b *billingCacheWorkerStub) SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *SubscriptionCacheData) error {
- atomic.AddInt64(&b.subscriptionUpdates, 1)
- return nil
- }
- func (b *billingCacheWorkerStub) UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error {
- atomic.AddInt64(&b.subscriptionUpdates, 1)
- return nil
- }
- func (b *billingCacheWorkerStub) InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error {
- return nil
- }
- func (b *billingCacheWorkerStub) GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*APIKeyRateLimitCacheData, error) {
- return nil, errors.New("not implemented")
- }
- func (b *billingCacheWorkerStub) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *APIKeyRateLimitCacheData) error {
- return nil
- }
- func (b *billingCacheWorkerStub) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
- return nil
- }
- func (b *billingCacheWorkerStub) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
- return nil
- }
- func (b *billingCacheWorkerStub) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaCacheEntry, bool, error) {
- return nil, false, nil
- }
- func (b *billingCacheWorkerStub) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error {
- return nil
- }
- func (b *billingCacheWorkerStub) DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error {
- return nil
- }
- func (b *billingCacheWorkerStub) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error {
- return nil
- }
- func (b *billingCacheWorkerStub) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]UserPlatformQuotaKey, error) {
- return nil, nil
- }
- func (b *billingCacheWorkerStub) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []UserPlatformQuotaKey) error {
- return nil
- }
- func (b *billingCacheWorkerStub) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []UserPlatformQuotaKey) ([]*UserPlatformQuotaCacheEntry, error) {
- return nil, nil
- }
- func TestBillingCacheServiceQueueHighLoad(t *testing.T) {
- cache := &billingCacheWorkerStub{}
- svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil)
- t.Cleanup(svc.Stop)
- start := time.Now()
- for i := 0; i < cacheWriteBufferSize*2; i++ {
- svc.QueueDeductBalance(1, 1)
- }
- require.Less(t, time.Since(start), 2*time.Second)
- svc.QueueUpdateSubscriptionUsage(1, 2, 1.5)
- require.Eventually(t, func() bool {
- return atomic.LoadInt64(&cache.balanceUpdates) > 0
- }, 2*time.Second, 10*time.Millisecond)
- require.Eventually(t, func() bool {
- return atomic.LoadInt64(&cache.subscriptionUpdates) > 0
- }, 2*time.Second, 10*time.Millisecond)
- }
- func TestBillingCacheServiceEnqueueAfterStopReturnsFalse(t *testing.T) {
- cache := &billingCacheWorkerStub{}
- svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, &config.Config{}, nil)
- svc.Stop()
- enqueued := svc.enqueueCacheWrite(cacheWriteTask{
- kind: cacheWriteDeductBalance,
- userID: 1,
- amount: 1,
- })
- require.False(t, enqueued)
- }
|