Files
sub2api/backend/internal/service/group_peak_rate_test.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

239 lines
8.5 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.
package service
import (
"context"
"math"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
)
func init() {
// 测试固定全局时区为 UTC,确保判定可复现。
_ = timezone.Init("UTC")
}
func newPeakGroup(enabled bool, start, end string, mult float64) *Group {
return &Group{
SubscriptionType: "subscription",
PeakRateEnabled: enabled,
PeakStart: start,
PeakEnd: end,
PeakRateMultiplier: mult,
}
}
func at(hour, min int) time.Time {
return time.Date(2026, 6, 29, hour, min, 0, 0, time.UTC)
}
func TestPeakMultiplierAt_DisabledOrUnconfigured(t *testing.T) {
cases := []struct {
name string
g *Group
}{
{"disabled", newPeakGroup(false, "14:00", "18:00", 3.0)},
{"empty start", newPeakGroup(true, "", "18:00", 3.0)},
{"empty end", newPeakGroup(true, "14:00", "", 3.0)},
{"invalid start>=end", newPeakGroup(true, "18:00", "14:00", 3.0)},
{"equal start==end", newPeakGroup(true, "14:00", "14:00", 3.0)},
{"malformed start", newPeakGroup(true, "99:99", "18:00", 3.0)},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
if got := c.g.PeakMultiplierAt(at(15, 0)); got != 1.0 {
t.Fatalf("expect 1.0, got %v", got)
}
})
}
}
func TestPeakMultiplierAt_NilReceiver(t *testing.T) {
var g *Group
if got := g.PeakMultiplierAt(at(15, 0)); got != 1.0 {
t.Fatalf("expect 1.0, got %v", got)
}
}
func TestPeakMultiplierAt_Boundaries(t *testing.T) {
g := newPeakGroup(true, "14:00", "18:00", 3.0)
cases := []struct {
t time.Time
want float64
}{
{at(13, 59), 1.0},
{at(14, 0), 3.0},
{at(15, 30), 3.0},
{at(17, 59), 3.0},
{at(18, 0), 1.0},
{at(23, 0), 1.0},
}
for _, c := range cases {
t.Run(c.t.Format("15:04"), func(t *testing.T) {
if got := g.PeakMultiplierAt(c.t); got != c.want {
t.Fatalf("at %s: expect %v, got %v", c.t.Format("15:04"), c.want, got)
}
})
}
}
func TestPeakMultiplierAt_RespectsTimezoneLocation(t *testing.T) {
// 全局时区为 UTC。北京 15:00 = UTC 07:00,不在 [14:00,18:00)。
nonUTC := time.Date(2026, 6, 29, 15, 0, 0, 0, mustLoad("Asia/Shanghai"))
g := newPeakGroup(true, "14:00", "18:00", 3.0)
if got := g.PeakMultiplierAt(nonUTC); got != 1.0 {
t.Fatalf("expect 1.0 (converted to UTC 07:00), got %v", got)
}
}
func mustLoad(name string) *time.Location {
loc, err := time.LoadLocation(name)
if err != nil {
panic(err)
}
return loc
}
func TestValidatePeakRateConfig(t *testing.T) {
cases := []struct {
name string
subType string
enabled bool
start string
end string
mult float64
wantErr bool
}{
{"disabled passes through", "subscription", false, "", "", 0, false},
{"subscription enabled valid", "subscription", true, "14:00", "18:00", 3.0, false},
{"standard enabled rejected", "standard", true, "14:00", "18:00", 3.0, true},
{"empty type treated as standard", "", true, "14:00", "18:00", 3.0, true},
{"standard disabled passes", "standard", false, "", "", 0, false},
{"enabled empty start", "subscription", true, "", "18:00", 1.0, true},
{"enabled empty end", "subscription", true, "14:00", "", 1.0, true},
{"enabled malformed start", "subscription", true, "99:99", "18:00", 1.0, true},
{"enabled malformed end", "subscription", true, "14:00", "25:00", 1.0, true},
{"enabled equal start==end", "subscription", true, "14:00", "14:00", 1.0, true},
{"enabled cross-day rejected", "subscription", true, "22:00", "02:00", 1.0, true},
{"enabled negative multiplier", "subscription", true, "14:00", "18:00", -0.5, true},
{"enabled zero multiplier allowed", "subscription", true, "14:00", "18:00", 0, false},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
err := ValidatePeakRateConfig(c.subType, c.enabled, c.start, c.end, c.mult)
if c.wantErr && err == nil {
t.Fatalf("expect error, got nil")
}
if !c.wantErr && err != nil {
t.Fatalf("expect no error, got %v", err)
}
})
}
}
func TestPeakMultiplierAt_StandardTypeDegradesToOne(t *testing.T) {
g := newPeakGroup(true, "14:00", "18:00", 3.0)
g.SubscriptionType = "standard"
if got := g.PeakMultiplierAt(at(15, 30)); got != 1.0 {
t.Fatalf("standard group must degrade to 1.0, got %v", got)
}
sub := newPeakGroup(true, "14:00", "18:00", 3.0)
sub.SubscriptionType = "subscription"
if got := sub.PeakMultiplierAt(at(15, 30)); got != 3.0 {
t.Fatalf("subscription group peak multiplier: got %v, want 3.0", got)
}
}
// TestPeakMultiplier_GatewayBillingSequence 调用 gateway_service.recordUsageCore 与
// openai_gateway_service.RecordUsage 共用的 computePeakAwareMultipliers,验证计费叠加顺序:
// 图片按次倍率基于基础倍率算出且不受高峰影响,高峰因子只乘入 token 倍率。
// 若有人调换叠加顺序或把高峰并入 imageMultiplier,此测试会失败。
func TestPeakMultiplier_GatewayBillingSequence(t *testing.T) {
const baseMultiplier = 0.8
apiKey := &APIKey{Group: newPeakGroup(true, "14:00", "18:00", 3.0)}
approxEq := func(a, b float64) bool { return math.Abs(a-b) < 1e-9 }
t.Run("peak hour amplifies token multiplier only", func(t *testing.T) {
now := at(15, 30) // 处于 [14:00, 18:00)
tokenMultiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, baseMultiplier, now)
if !approxEq(imageMultiplier, baseMultiplier) {
t.Fatalf("image multiplier must not be affected by peak: got %v, want %v", imageMultiplier, baseMultiplier)
}
if want := baseMultiplier * 3.0; !approxEq(tokenMultiplier, want) {
t.Fatalf("token multiplier should include peak factor: got %v, want %v", tokenMultiplier, want)
}
})
t.Run("off-peak leaves both multipliers at base", func(t *testing.T) {
now := at(20, 0)
tokenMultiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, baseMultiplier, now)
if !approxEq(imageMultiplier, baseMultiplier) {
t.Fatalf("image multiplier: got %v, want %v", imageMultiplier, baseMultiplier)
}
if !approxEq(tokenMultiplier, baseMultiplier) {
t.Fatalf("token multiplier should equal base off-peak: got %v, want %v", tokenMultiplier, baseMultiplier)
}
})
t.Run("image independent mode decoupled from peak", func(t *testing.T) {
indGroup := newPeakGroup(true, "14:00", "18:00", 3.0)
indGroup.ImageRateIndependent = true
indGroup.ImageRateMultiplier = 0.5
indKey := &APIKey{Group: indGroup}
now := at(15, 30)
tokenMultiplier, imageMultiplier := computePeakAwareMultipliers(indKey, baseMultiplier, now)
if !approxEq(imageMultiplier, 0.5) {
t.Fatalf("independent image multiplier: got %v, want 0.5", imageMultiplier)
}
if want := baseMultiplier * 3.0; !approxEq(tokenMultiplier, want) {
t.Fatalf("token multiplier should include peak factor: got %v, want %v", tokenMultiplier, want)
}
})
t.Run("nil api key degrades to base multipliers", func(t *testing.T) {
now := at(15, 30)
tokenMultiplier, imageMultiplier := computePeakAwareMultipliers(nil, baseMultiplier, now)
if !approxEq(tokenMultiplier, baseMultiplier) {
t.Fatalf("nil group token multiplier: got %v, want %v", tokenMultiplier, baseMultiplier)
}
if !approxEq(imageMultiplier, baseMultiplier) {
t.Fatalf("nil group image multiplier: got %v, want %v", imageMultiplier, baseMultiplier)
}
})
}
// TestPeakMultiplier_SnapshotRoundTrip 防回归:认证缓存快照(APIKeyAuthGroupSnapshot
// 必须携带高峰倍率 4 字段,否则扣费路径拿到的 apiKey.Group 会缺字段、PeakMultiplierAt 恒降级为 1.0。
// 调用真实链路 snapshotFromAPIKey → snapshotToAPIKey,验证 peak 配置经快照往返后仍生效。
func TestPeakMultiplier_SnapshotRoundTrip(t *testing.T) {
apiKey := &APIKey{
User: &User{ID: 1, Status: StatusActive, Role: RoleUser},
Group: newPeakGroup(true, "14:00", "18:00", 3.0),
}
svc := &APIKeyService{}
snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
if snapshot == nil || snapshot.Group == nil {
t.Fatalf("snapshot or snapshot.Group must not be nil")
}
restored := svc.snapshotToAPIKey("k", snapshot)
if restored.Group == nil {
t.Fatalf("restored.Group must not be nil")
}
if !restored.Group.PeakRateEnabled ||
restored.Group.PeakStart != "14:00" ||
restored.Group.PeakEnd != "18:00" ||
restored.Group.PeakRateMultiplier != 3.0 {
t.Fatalf("peak fields lost in snapshot round-trip: %+v", restored.Group)
}
if got := restored.Group.PeakMultiplierAt(at(15, 30)); got != 3.0 {
t.Fatalf("peak hour multiplier after round-trip: got %v, want 3.0", got)
}
if got := restored.Group.PeakMultiplierAt(at(20, 0)); got != 1.0 {
t.Fatalf("off-peak multiplier after round-trip: got %v, want 1.0", got)
}
}