Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/billing_cache_service_balance_test.go

@@ -0,0 +1,128 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/stretchr/testify/require"
+)
+
+type balanceEligibilityCacheStub struct {
+	billingCacheWorkerStub
+
+	balance                  float64
+	cacheMissAfterInvalidate bool
+	invalidated              atomic.Bool
+	deductCalls              atomic.Int64
+	invalidateCalls          atomic.Int64
+}
+
+func (s *balanceEligibilityCacheStub) GetUserBalance(context.Context, int64) (float64, error) {
+	if s.cacheMissAfterInvalidate && s.invalidated.Load() {
+		return 0, errors.New("cache miss")
+	}
+	return s.balance, nil
+}
+
+func (s *balanceEligibilityCacheStub) DeductUserBalance(context.Context, int64, float64) error {
+	s.deductCalls.Add(1)
+	return nil
+}
+
+func (s *balanceEligibilityCacheStub) InvalidateUserBalance(context.Context, int64) error {
+	s.invalidateCalls.Add(1)
+	s.invalidated.Store(true)
+	return nil
+}
+
+func TestCheckBillingEligibility_RejectsBalanceBelowMinimumReserve(t *testing.T) {
+	cache := &balanceEligibilityCacheStub{balance: 0.005}
+	cfg := &config.Config{}
+	cfg.Billing.MinimumBalanceReserve = 0.01
+	svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
+	t.Cleanup(svc.Stop)
+
+	err := svc.CheckBillingEligibility(context.Background(), &User{ID: 1}, nil, nil, nil, "")
+	require.ErrorIs(t, err, ErrInsufficientBalance)
+}
+
+func TestCheckBillingEligibility_AllowsBalanceAtMinimumReserve(t *testing.T) {
+	cache := &balanceEligibilityCacheStub{balance: 0.01}
+	cfg := &config.Config{}
+	cfg.Billing.MinimumBalanceReserve = 0.01
+	svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
+	t.Cleanup(svc.Stop)
+
+	err := svc.CheckBillingEligibility(context.Background(), &User{ID: 1}, nil, nil, nil, "")
+	require.NoError(t, err)
+}
+
+func TestSyncBalanceCacheAfterDeduction_InvalidatesExhaustedBalance(t *testing.T) {
+	cache := &balanceEligibilityCacheStub{
+		balance:                  0.50,
+		cacheMissAfterInvalidate: true,
+	}
+	userRepo := &balanceLoadUserRepoStub{balance: -0.25}
+	cfg := &config.Config{}
+	cfg.Billing.MinimumBalanceReserve = 0.01
+	svc := NewBillingCacheService(cache, userRepo, nil, nil, nil, nil, cfg, nil)
+	t.Cleanup(svc.Stop)
+
+	newBalance := -0.25
+	syncBalanceCacheAfterDeduction(context.Background(), &postUsageBillingParams{
+		Cost: &CostBreakdown{ActualCost: 0.75},
+		User: &User{ID: 1},
+	}, &billingDeps{billingCacheService: svc}, &UsageBillingApplyResult{
+		NewBalance:         &newBalance,
+		BalanceOverdrafted: true,
+	})
+
+	require.Equal(t, int64(1), cache.invalidateCalls.Load())
+	require.Equal(t, int64(0), cache.deductCalls.Load())
+
+	err := svc.CheckBillingEligibility(context.Background(), &User{ID: 1}, nil, nil, nil, "")
+	require.ErrorIs(t, err, ErrInsufficientBalance)
+	require.Equal(t, int64(1), userRepo.calls.Load())
+}
+
+func TestSyncBalanceCacheAfterDeduction_InvalidatesWhenBalanceFallsBelowReserve(t *testing.T) {
+	cache := &balanceEligibilityCacheStub{balance: 0.50}
+	cfg := &config.Config{}
+	cfg.Billing.MinimumBalanceReserve = 0.01
+	svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
+	t.Cleanup(svc.Stop)
+
+	newBalance := 0.005
+	syncBalanceCacheAfterDeduction(context.Background(), &postUsageBillingParams{
+		Cost: &CostBreakdown{ActualCost: 0.495},
+		User: &User{ID: 1},
+	}, &billingDeps{billingCacheService: svc}, &UsageBillingApplyResult{NewBalance: &newBalance})
+
+	require.Equal(t, int64(1), cache.invalidateCalls.Load())
+	require.Equal(t, int64(0), cache.deductCalls.Load())
+}
+
+func TestSyncBalanceCacheAfterDeduction_QueuesDeductWhenBalanceStillEligible(t *testing.T) {
+	cache := &balanceEligibilityCacheStub{balance: 1}
+	cfg := &config.Config{}
+	cfg.Billing.MinimumBalanceReserve = 0.01
+	svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
+	t.Cleanup(svc.Stop)
+
+	newBalance := 0.75
+	syncBalanceCacheAfterDeduction(context.Background(), &postUsageBillingParams{
+		Cost: &CostBreakdown{ActualCost: 0.25},
+		User: &User{ID: 1},
+	}, &billingDeps{billingCacheService: svc}, &UsageBillingApplyResult{NewBalance: &newBalance})
+
+	require.Equal(t, int64(0), cache.invalidateCalls.Load())
+	require.Eventually(t, func() bool {
+		return cache.deductCalls.Load() == 1
+	}, 2*time.Second, 10*time.Millisecond)
+}