6b39b344d8
为用户在 anthropic/openai/gemini/antigravity 四个平台上提供日/周/月 三个窗口的 USD 配额管控。配额语义：未设置=不限制，0=禁用，>0=美元上限。 两层模型： - 配置层：系统默认配额，以及 email/linuxdo/oidc/wechat/github/google/ dingtalk 七个鉴权来源的默认配额，存于 settings，以嵌套 JSON 整体读写 （系统 1 个 key + 每个来源 1 个 key），整体替换语义。 - 运行时层：user_platform_quota 表按用户记录实际配额，与配置层解耦。 后端：新增 ent schema 与 140_user_platform_quotas.sql 迁移、repository 与 service 端口、计费链路集成、管理端与用户端读写接口。 前端：管理端设置页配额编辑、用户配额管理 Modal、用户 Dashboard 展示、 中英文案。 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
121 lines
3.7 KiB
Go
121 lines
3.7 KiB
Go
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) error {
|
|
return 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)
|
|
}
|