为用户在 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>
158 lines
5.4 KiB
Go
158 lines
5.4 KiB
Go
//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 早退路径仍携带全局 quota(GetDefaultPlatformQuotas 在 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")
|
||
}
|
||
}
|