mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
239 lines
12 KiB
Go
239 lines
12 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestOpenAISubscriptionLongContextExempt(t *testing.T) {
|
|
oauth := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
for _, tt := range []struct {
|
|
name string
|
|
account *Account
|
|
result *OpenAIForwardResult
|
|
want bool
|
|
}{
|
|
{name: "空账号", result: &OpenAIForwardResult{Model: "gpt-6.1-sol"}},
|
|
{name: "空结果", account: oauth},
|
|
{name: "订阅", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6.1-sol"}, want: true},
|
|
{name: "别名", account: oauth, result: &OpenAIForwardResult{Model: "openai/GPT_6.1_SOL-max"}, want: true},
|
|
{name: "压缩", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6.1-sol-openai-compact"}, want: true},
|
|
{name: "映射到目标", account: oauth, result: &OpenAIForwardResult{Model: "custom", UpstreamModel: "gpt-6.1-sol"}, want: true},
|
|
{name: "映射离开目标", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6.1-sol", UpstreamModel: "gpt-6-sol"}},
|
|
{name: "仅计费标签相同", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6-sol", BillingModel: "gpt-6.1-sol"}},
|
|
{name: "仅响应标签相同", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6-sol", UpstreamResponseModel: "gpt-6.1-sol"}},
|
|
{name: "其他模型", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6-astra"}},
|
|
{name: "相似名称", account: oauth, result: &OpenAIForwardResult{Model: "gpt-6.1-sol-preview"}},
|
|
{name: "API账号", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, result: &OpenAIForwardResult{Model: "gpt-6.1-sol"}},
|
|
{name: "其他平台", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, result: &OpenAIForwardResult{Model: "gpt-6.1-sol"}},
|
|
{name: "订阅访问令牌", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeSetupToken}, result: &OpenAIForwardResult{Model: "gpt-6.1-sol"}, want: true},
|
|
} {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
require.Equal(t, tt.want, openAISubscriptionLongContextExempt(tt.account, tt.result))
|
|
})
|
|
}
|
|
}
|
|
|
|
// 同时覆盖远程目录更新、静态兜底、旧计费入口和分组入口，校验真实扣款与账号成本。
|
|
func TestOpenAIGatewayServiceRecordUsage_GPT61SolSubscriptionLongContext(t *testing.T) {
|
|
for _, catalog := range []bool{false, true} {
|
|
for _, groupMode := range []string{"无分组", "分组关闭", "分组开启"} {
|
|
for _, accountEnabled := range []bool{false, true} {
|
|
for _, accountType := range []string{AccountTypeOAuth, AccountTypeSetupToken, AccountTypeAPIKey} {
|
|
for _, tier := range []string{"", "priority"} {
|
|
name := fmt.Sprintf("目录=%t/%s/账号开关=%t/%s/%s", catalog, groupMode, accountEnabled, accountType, tier)
|
|
t.Run(name, func(t *testing.T) {
|
|
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
|
userRepo := &openAIRecordUsageUserRepoStub{}
|
|
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{}, nil)
|
|
if catalog {
|
|
svc.billingService = NewBillingService(svc.cfg, newStubPricingServiceFromJSON(t, `{
|
|
"gpt-6.1-sol": {"litellm_provider": "openai", "mode": "chat",
|
|
"input_cost_per_token": 2e-6, "output_cost_per_token": 10e-6,
|
|
"cache_read_input_token_cost": 0.1e-6, "cache_creation_input_token_cost": 2.5e-6,
|
|
"input_cost_per_token_above_272k_tokens": 4e-6,
|
|
"output_cost_per_token_above_272k_tokens": 15e-6}
|
|
}`))
|
|
}
|
|
apiKey := &APIKey{ID: 1}
|
|
if groupMode != "无分组" {
|
|
apiKey = openAIRecordUsageAPIKeyWithGroup(svc, 1, groupMode == "分组开启")
|
|
apiKey.GroupID = i64p(1)
|
|
apiKey.Group.RateMultiplier = 1.1
|
|
svc.channelService = newTestChannelServiceForStats(t, &Channel{ID: 1, Status: StatusActive}, 1, PlatformOpenAI)
|
|
}
|
|
account := &Account{ID: 3, Platform: PlatformOpenAI, Type: accountType,
|
|
Extra: map[string]any{openAILongContextBillingEnabledKey: accountEnabled}}
|
|
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
|
Result: &OpenAIForwardResult{
|
|
RequestID: name, Model: "gpt-6.1-sol", ServiceTier: &tier, Duration: time.Second,
|
|
// 普通输入只有 200k，由缓存使总上下文越过 272k。
|
|
Usage: OpenAIUsage{InputTokens: 300000, OutputTokens: 2000, CacheReadInputTokens: 80000, CacheCreationInputTokens: 20000},
|
|
},
|
|
APIKey: apiKey, User: &User{ID: 2}, Account: account,
|
|
})
|
|
require.NoError(t, err)
|
|
wantLong := accountType == AccountTypeAPIKey && (accountEnabled || groupMode == "分组开启")
|
|
inputFactor, outputFactor, tierFactor := 1.0, 1.0, 1.0
|
|
if wantLong {
|
|
inputFactor, outputFactor = 2, 1.5
|
|
}
|
|
if tier == "priority" {
|
|
tierFactor = 2
|
|
}
|
|
log := usageRepo.lastLog
|
|
require.NotNil(t, log)
|
|
require.Equal(t, wantLong, log.LongContextBillingApplied)
|
|
require.InDelta(t, 200000*2e-6*inputFactor*tierFactor, log.InputCost, 1e-12)
|
|
require.InDelta(t, 80000*0.1e-6*inputFactor*tierFactor, log.CacheReadCost, 1e-12)
|
|
require.InDelta(t, 20000*2.5e-6*inputFactor*tierFactor, log.CacheCreationCost, 1e-12)
|
|
require.InDelta(t, 2000*10e-6*outputFactor*tierFactor, log.OutputCost, 1e-12)
|
|
wantTotal := (0.458*inputFactor + 0.02*outputFactor) * tierFactor
|
|
require.InDelta(t, wantTotal, log.TotalCost, 1e-12)
|
|
require.InDelta(t, wantTotal*1.1, log.ActualCost, 1e-12)
|
|
require.Equal(t, 1, userRepo.deductCalls)
|
|
require.InDelta(t, wantTotal*1.1, userRepo.lastAmount, 1e-12)
|
|
if groupMode != "无分组" {
|
|
wantStats := 0.478 * tierFactor
|
|
if accountType == AccountTypeAPIKey && accountEnabled {
|
|
wantStats = (0.458*2 + 0.02*1.5) * tierFactor
|
|
}
|
|
require.NotNil(t, log.AccountStatsCost)
|
|
require.InDelta(t, wantStats, *log.AccountStatsCost, 1e-12)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceRecordUsage_GPT61SolUsesCredentialAndUpstreamModel(t *testing.T) {
|
|
for _, tt := range []struct {
|
|
name string
|
|
model string
|
|
upstream string
|
|
parentType string
|
|
wantTotal float64
|
|
wantLong bool
|
|
}{
|
|
{name: "客户别名映射到目标", model: "custom", upstream: "gpt-6.1-sol", parentType: AccountTypeOAuth, wantTotal: 0.62},
|
|
{name: "请求目标映射到旧模型", model: "gpt-6.1-sol", upstream: "gpt-6-sol", parentType: AccountTypeOAuth, wantTotal: 1.23, wantLong: true},
|
|
{name: "其他订阅模型", model: "gpt-6-sol", parentType: AccountTypeOAuth, wantTotal: 1.23, wantLong: true},
|
|
{name: "订阅母账号", model: "gpt-6.1-sol", parentType: AccountTypeOAuth, wantTotal: 0.62},
|
|
} {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
|
svc := newOpenAIRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{}, nil)
|
|
parentID := int64(4)
|
|
svc.accountRepo = &openAIRecordUsageAccountRepoStub{account: &Account{
|
|
ID: parentID, Platform: PlatformOpenAI, Type: tt.parentType,
|
|
Extra: map[string]any{openAILongContextBillingEnabledKey: true},
|
|
}}
|
|
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
|
Result: &OpenAIForwardResult{RequestID: tt.name, Model: tt.model, UpstreamModel: tt.upstream,
|
|
Usage: OpenAIUsage{InputTokens: 300000, OutputTokens: 2000}, Duration: time.Second},
|
|
APIKey: openAIRecordUsageAPIKeyWithGroup(svc, 1, true), User: &User{ID: 2},
|
|
Account: &Account{ID: 3, Platform: PlatformOpenAI, Type: AccountTypeOAuth, ParentAccountID: &parentID, QuotaDimension: QuotaDimensionSpark},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, tt.wantLong, usageRepo.lastLog.LongContextBillingApplied)
|
|
require.InDelta(t, tt.wantTotal, usageRepo.lastLog.TotalCost, 1e-12)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCalculateCostUnified_ModelLongContextExemptionPreservesCustomIntervals(t *testing.T) {
|
|
service := NewBillingService(nil, nil)
|
|
for _, withResolver := range []bool{false, true} {
|
|
input := CostInput{Model: "gpt-6.1-sol", Tokens: UsageTokens{InputTokens: 300000}, RateMultiplier: 1,
|
|
LongContextBillingEnabled: boolPtr(true), DisableModelLongContextPricing: true}
|
|
if withResolver {
|
|
input.Ctx = context.Background()
|
|
input.Resolver = NewModelPricingResolver(nil, service)
|
|
}
|
|
cost, err := service.CalculateCostUnified(input)
|
|
require.NoError(t, err)
|
|
require.InDelta(t, 0.6, cost.TotalCost, 1e-12)
|
|
require.False(t, cost.LongContextBillingApplied)
|
|
if withResolver {
|
|
input.Resolved = &ResolvedPricing{
|
|
BasePricing: &ModelPricing{InputPricePerToken: 2e-6},
|
|
Intervals: []PricingInterval{{MinTokens: 272000, InputMultiplier: pricingMultiplier(3)}},
|
|
longContextPricingEnabled: true,
|
|
}
|
|
cost, err = service.CalculateCostUnified(input)
|
|
require.NoError(t, err)
|
|
require.InDelta(t, 1.8, cost.TotalCost, 1e-12, "管理员显式区间价格应保留")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceRecordUsage_GPT61SolLongContextBoundary(t *testing.T) {
|
|
for _, totalInput := range []int{271999, 272000, 272001} {
|
|
for _, accountType := range []string{AccountTypeOAuth, AccountTypeAPIKey} {
|
|
t.Run(fmt.Sprintf("%s/%d", accountType, totalInput), func(t *testing.T) {
|
|
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
|
svc := newOpenAIRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{}, nil)
|
|
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
|
Result: &OpenAIForwardResult{RequestID: "boundary", Model: "gpt-6.1-sol", Duration: time.Second,
|
|
Usage: OpenAIUsage{InputTokens: totalInput, OutputTokens: 1000}},
|
|
APIKey: openAIRecordUsageAPIKeyWithGroup(svc, 1, true), User: &User{ID: 2},
|
|
Account: &Account{ID: 3, Platform: PlatformOpenAI, Type: accountType,
|
|
Extra: map[string]any{openAILongContextBillingEnabledKey: true}},
|
|
})
|
|
require.NoError(t, err)
|
|
wantLong := accountType == AccountTypeAPIKey && totalInput > 272000
|
|
wantInput, wantOutput := float64(totalInput)*2e-6, 0.01
|
|
if wantLong {
|
|
wantInput *= 2
|
|
wantOutput *= 1.5
|
|
}
|
|
require.Equal(t, wantLong, usageRepo.lastLog.LongContextBillingApplied)
|
|
require.InDelta(t, wantInput, usageRepo.lastLog.InputCost, 1e-12)
|
|
require.InDelta(t, wantOutput, usageRepo.lastLog.OutputCost, 1e-12)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceRecordUsage_GPT61SolExemptionSurvivesRepricing(t *testing.T) {
|
|
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
|
userRepo := &openAIRecordUsageUserRepoStub{}
|
|
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, &openAIRecordUsageSubRepoStub{}, nil)
|
|
apiKey := openAIRecordUsageAPIKeyWithGroup(svc, 1, true)
|
|
apiKey.GroupID = i64p(1)
|
|
apiKey.Group.Platform = PlatformOpenAI
|
|
apiKey.Group.RateMultiplier = 1.1
|
|
apiKey.Group.FreeOpenAIFast = true
|
|
svc.channelService = newTestChannelServiceForStats(t, &Channel{ID: 1, Status: StatusActive}, 1, PlatformOpenAI)
|
|
tier := "priority"
|
|
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
|
Result: &OpenAIForwardResult{
|
|
RequestID: "repricing", Model: "gpt-6-sol", UpstreamModel: "gpt-6.1-sol", BillingModel: "gpt-6-sol",
|
|
UpstreamResponseModel: "gpt-6.1-sol", ServiceTier: &tier, Duration: time.Second,
|
|
Usage: OpenAIUsage{InputTokens: 300000, OutputTokens: 2000, CacheReadInputTokens: 80000},
|
|
},
|
|
APIKey: apiKey, User: &User{ID: 2},
|
|
Account: &Account{ID: 3, Platform: PlatformOpenAI, Type: AccountTypeOAuth,
|
|
Extra: map[string]any{openAILongContextBillingEnabledKey: true}},
|
|
ChannelUsageFields: ChannelUsageFields{BillingModelSource: BillingModelSourceResponse},
|
|
})
|
|
require.NoError(t, err)
|
|
// 响应模型改用较低的缓存价；免费 Fast 再按标准档扣款，两次重算均须保留豁免。
|
|
const standardCost = 220000*2e-6 + 80000*0.1e-6 + 2000*10e-6
|
|
require.InDelta(t, standardCost*2, usageRepo.lastLog.TotalCost, 1e-12)
|
|
require.InDelta(t, standardCost*1.1, usageRepo.lastLog.ActualCost, 1e-12)
|
|
require.InDelta(t, standardCost*1.1, userRepo.lastAmount, 1e-12)
|
|
require.False(t, usageRepo.lastLog.LongContextBillingApplied)
|
|
require.NotNil(t, usageRepo.lastLog.AccountStatsCost)
|
|
require.InDelta(t, standardCost*2, *usageRepo.lastLog.AccountStatsCost, 1e-12)
|
|
}
|