sub2api/backend/internal/service/auth_service_platform_quota_test.go
DaydreamCoding 6b39b344d8 feat(quota): 用户 × 平台 USD 配额
为用户在 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>
2026-05-26 10:49:20 +08:00

158 lines
5.4 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//go:build unit
package service
import (
"context"
"fmt"
"testing"
"time"
)
// fakeInsertRecorder 记录 BulkInsertInitial 调用,实现 UserPlatformQuotaRepository port。
type fakeInsertRecorder struct {
records []UserPlatformQuotaRecord
err error
}
func (f *fakeInsertRecorder) GetByUserPlatform(_ context.Context, _ int64, _ string) (*UserPlatformQuotaRecord, error) {
return nil, nil
}
func (f *fakeInsertRecorder) BulkInsertInitial(_ context.Context, recs []UserPlatformQuotaRecord) error {
if f.err != nil {
return f.err
}
f.records = append(f.records, recs...)
return nil
}
func (f *fakeInsertRecorder) IncrementUsageWithReset(_ context.Context, _ int64, _ string, _ float64, _ time.Time) error {
return nil
}
func (f *fakeInsertRecorder) ListByUser(_ context.Context, _ int64) ([]UserPlatformQuotaRecord, error) {
return nil, nil
}
func (f *fakeInsertRecorder) UpsertForUser(_ context.Context, _ int64, _ []UserPlatformQuotaRecord) error {
return nil
}
func (f *fakeInsertRecorder) ResetExpiredWindow(_ context.Context, _ int64, _ string, _ string, _ time.Time) error {
return nil
}
func TestSnapshotPlatformQuotaDefaults_PassesToRepoBulkInsert(t *testing.T) {
fakeRepo := &fakeInsertRecorder{}
s := &AuthService{userPlatformQuotaRepo: fakeRepo}
five := 5.0
plan := &signupGrantPlan{
PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{
"anthropic": {DailyLimitUSD: &five},
"openai": {},
"gemini": {},
"antigravity": {},
},
}
if err := s.snapshotPlatformQuotaDefaults(context.Background(), 999, plan); err != nil {
t.Fatal(err)
}
if len(fakeRepo.records) != 4 {
t.Errorf("expected 4 records, got %d", len(fakeRepo.records))
}
found := false
for _, r := range fakeRepo.records {
if r.UserID == 999 && r.Platform == "anthropic" && r.DailyLimitUSD != nil && *r.DailyLimitUSD == 5 {
found = true
}
}
if !found {
t.Error("anthropic daily = 5 not snapshotted")
}
}
func TestSnapshotPlatformQuotaDefaults_NilPlanIsNoop(t *testing.T) {
fakeRepo := &fakeInsertRecorder{}
s := &AuthService{userPlatformQuotaRepo: fakeRepo}
if err := s.snapshotPlatformQuotaDefaults(context.Background(), 1, nil); err != nil {
t.Errorf("nil plan should be noop, got %v", err)
}
if len(fakeRepo.records) != 0 {
t.Errorf("expected no records, got %d", len(fakeRepo.records))
}
}
func TestSnapshotPlatformQuotaDefaults_RepoErrorFailsOpen(t *testing.T) {
fakeRepo := &fakeInsertRecorder{err: fmt.Errorf("db down")}
s := &AuthService{userPlatformQuotaRepo: fakeRepo}
five := 5.0
plan := &signupGrantPlan{
PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{
"anthropic": {DailyLimitUSD: &five},
},
}
if err := s.snapshotPlatformQuotaDefaults(context.Background(), 1, plan); err != nil {
t.Errorf("fail-open: expected nil even on repo error, got %v", err)
}
}
func TestSnapshotPlatformQuotaDefaults_NilRepoIsNoop(t *testing.T) {
s := &AuthService{userPlatformQuotaRepo: nil}
five := 5.0
plan := &signupGrantPlan{
PlatformQuotas: map[string]*DefaultPlatformQuotaSetting{"a": {DailyLimitUSD: &five}},
}
if err := s.snapshotPlatformQuotaDefaults(context.Background(), 1, plan); err != nil {
t.Errorf("nil repo should be noop, got %v", err)
}
}
// resolveSignupGrantPlan 测试:依赖完整的 AuthService 构造,需要 SettingService含 settingRepoStub
// settingRepoStub 已在 auth_service_register_test.go 中定义,同 package 可直接使用。
func TestResolveSignupGrantPlan_GlobalQuotaLoadedBeforeAuthSource(t *testing.T) {
// 全局 quota JSON key新格式
settings := map[string]string{
SettingKeyRegistrationEnabled: "true",
SettingKeyDefaultPlatformQuotas: `{
"anthropic": {"daily": 10, "weekly": 50, "monthly": 200},
"openai": {"daily": 5, "weekly": 25, "monthly": 100},
"gemini": {"daily": 5, "weekly": 25, "monthly": 100},
"antigravity": {"daily": 5, "weekly": 25, "monthly": 100}
}`,
}
svc := newAuthService(nil, settings, nil, nil)
plan := svc.resolveSignupGrantPlan(context.Background(), "email")
if plan.PlatformQuotas == nil {
t.Fatal("expected PlatformQuotas to be non-nil after loading global quota KVs")
}
q := plan.PlatformQuotas["anthropic"]
if q == nil {
t.Fatal("expected anthropic quota to be set")
}
if q.DailyLimitUSD == nil || *q.DailyLimitUSD != 10 {
t.Errorf("expected anthropic daily=10, got %v", q.DailyLimitUSD)
}
}
// TestResolveSignupGrantPlan_DisabledAuthSourceStillCarriesGlobalQuota 验证 P1 约束:
// !enabled 早退路径仍携带全局 quotaGetDefaultPlatformQuotas 在 ResolveAuthSourceGrantSettings 之前)。
func TestResolveSignupGrantPlan_DisabledAuthSourceStillCarriesGlobalQuota(t *testing.T) {
settings := map[string]string{
SettingKeyRegistrationEnabled: "true",
// auth source 不配置(=> !enabled 路径)
SettingKeyDefaultPlatformQuotas: `{"anthropic": {"daily": 10, "weekly": 50, "monthly": 200}}`,
}
svc := newAuthService(nil, settings, nil, nil)
plan := svc.resolveSignupGrantPlan(context.Background(), "email")
// !enabled 路径plan.PlatformQuotas 应已含全局层(不是 nil
if plan.PlatformQuotas == nil {
t.Fatal("P1 violated: PlatformQuotas is nil even with global quota KVs set")
}
// P1 核心断言disabled auth source 路径不能丢失全局 quota
if _, ok := plan.PlatformQuotas["anthropic"]; !ok {
t.Error("P1 violated: disabled auth source path dropped global platform quota")
}
}