| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057 |
- //go:build unit
- package service
- import (
- "context"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/stretchr/testify/require"
- )
- // ---------------------------------------------------------------------------
- // matchAccountStatsRule
- // ---------------------------------------------------------------------------
- func TestMatchAccountStatsRule_BothEmpty_NoMatch(t *testing.T) {
- rule := &AccountStatsPricingRule{}
- require.False(t, matchAccountStatsRule(rule, 1, 10))
- }
- func TestMatchAccountStatsRule_AccountIDMatch(t *testing.T) {
- rule := &AccountStatsPricingRule{AccountIDs: []int64{1, 2, 3}}
- require.True(t, matchAccountStatsRule(rule, 2, 999))
- }
- func TestMatchAccountStatsRule_GroupIDMatch(t *testing.T) {
- rule := &AccountStatsPricingRule{GroupIDs: []int64{10, 20}}
- require.True(t, matchAccountStatsRule(rule, 999, 20))
- }
- func TestMatchAccountStatsRule_BothConfigured_AccountMatch(t *testing.T) {
- rule := &AccountStatsPricingRule{
- AccountIDs: []int64{1, 2},
- GroupIDs: []int64{10, 20},
- }
- require.True(t, matchAccountStatsRule(rule, 2, 999))
- }
- func TestMatchAccountStatsRule_BothConfigured_GroupMatch(t *testing.T) {
- rule := &AccountStatsPricingRule{
- AccountIDs: []int64{1, 2},
- GroupIDs: []int64{10, 20},
- }
- require.True(t, matchAccountStatsRule(rule, 999, 10))
- }
- func TestMatchAccountStatsRule_BothConfigured_NeitherMatch(t *testing.T) {
- rule := &AccountStatsPricingRule{
- AccountIDs: []int64{1, 2},
- GroupIDs: []int64{10, 20},
- }
- require.False(t, matchAccountStatsRule(rule, 999, 999))
- }
- // ---------------------------------------------------------------------------
- // findPricingForModel
- // ---------------------------------------------------------------------------
- func TestFindPricingForModel(t *testing.T) {
- exactPricing := ChannelModelPricing{
- ID: 1,
- Models: []string{"claude-opus-4"},
- }
- wildcardPricing := ChannelModelPricing{
- ID: 2,
- Models: []string{"claude-*"},
- }
- platformPricing := ChannelModelPricing{
- ID: 3,
- Platform: "openai",
- Models: []string{"gpt-4o"},
- }
- emptyPlatformPricing := ChannelModelPricing{
- ID: 4,
- Models: []string{"gemini-2.5-pro"},
- }
- tests := []struct {
- name string
- list []ChannelModelPricing
- platform string
- model string
- wantID int64
- wantNil bool
- }{
- {
- name: "exact match",
- list: []ChannelModelPricing{exactPricing},
- platform: "anthropic",
- model: "claude-opus-4",
- wantID: 1,
- },
- {
- name: "exact match case insensitive",
- list: []ChannelModelPricing{{ID: 5, Models: []string{"Claude-Opus-4"}}},
- platform: "",
- model: "claude-opus-4",
- wantID: 5,
- },
- {
- name: "wildcard match",
- list: []ChannelModelPricing{wildcardPricing},
- platform: "anthropic",
- model: "claude-opus-4",
- wantID: 2,
- },
- {
- name: "exact match takes priority over wildcard",
- list: []ChannelModelPricing{wildcardPricing, exactPricing},
- platform: "anthropic",
- model: "claude-opus-4",
- wantID: 1,
- },
- {
- name: "platform mismatch skipped",
- list: []ChannelModelPricing{platformPricing},
- platform: "anthropic",
- model: "gpt-4o",
- wantNil: true,
- },
- {
- name: "empty platform in pricing matches any",
- list: []ChannelModelPricing{emptyPlatformPricing},
- platform: "gemini",
- model: "gemini-2.5-pro",
- wantID: 4,
- },
- {
- name: "empty platform in query matches any pricing platform",
- list: []ChannelModelPricing{platformPricing},
- platform: "",
- model: "gpt-4o",
- wantID: 3,
- },
- {
- name: "no match at all",
- list: []ChannelModelPricing{exactPricing, wildcardPricing},
- platform: "anthropic",
- model: "gpt-4o",
- wantNil: true,
- },
- {
- name: "empty list returns nil",
- list: nil,
- model: "claude-opus-4",
- wantNil: true,
- },
- {
- name: "wildcard matches by config order (first match wins)",
- list: []ChannelModelPricing{
- {ID: 10, Models: []string{"claude-*"}},
- {ID: 11, Models: []string{"claude-opus-*"}},
- },
- platform: "",
- model: "claude-opus-4",
- wantID: 10, // config order: "claude-*" is first and matches, so it wins
- },
- {
- name: "shorter wildcard used when longer does not match",
- list: []ChannelModelPricing{
- {ID: 10, Models: []string{"claude-*"}},
- {ID: 11, Models: []string{"claude-opus-*"}},
- },
- platform: "",
- model: "claude-sonnet-4",
- wantID: 10, // only "claude-*" matches
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- result := findPricingForModel(tt.list, tt.platform, tt.model)
- if tt.wantNil {
- require.Nil(t, result)
- return
- }
- require.NotNil(t, result)
- require.Equal(t, tt.wantID, result.ID)
- })
- }
- }
- // ---------------------------------------------------------------------------
- // calculateStatsCost
- // ---------------------------------------------------------------------------
- func TestCalculateStatsCost_NilPricing(t *testing.T) {
- result := calculateStatsCost(nil, UsageTokens{}, 1)
- require.Nil(t, result)
- }
- func TestCalculateStatsCost_TokenBilling(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeToken,
- InputPrice: testPtrFloat64(0.001),
- OutputPrice: testPtrFloat64(0.002),
- }
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- }
- result := calculateStatsCost(pricing, tokens, 1)
- require.NotNil(t, result)
- // 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
- require.InDelta(t, 0.2, *result, 1e-12)
- }
- func TestCalculateStatsCost_TokenBilling_WithCache(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeToken,
- InputPrice: testPtrFloat64(0.001),
- OutputPrice: testPtrFloat64(0.002),
- CacheWritePrice: testPtrFloat64(0.003),
- CacheReadPrice: testPtrFloat64(0.0005),
- }
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- CacheCreationTokens: 200,
- CacheReadTokens: 300,
- }
- result := calculateStatsCost(pricing, tokens, 1)
- require.NotNil(t, result)
- // 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
- // = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
- require.InDelta(t, 0.95, *result, 1e-12)
- }
- func TestCalculateStatsCost_TokenBilling_WithCacheTTLPrices(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeToken,
- CacheWritePrice: testPtrFloat64(0.003),
- CacheWrite1hPrice: testPtrFloat64(0.005),
- }
- tokens := UsageTokens{
- CacheCreationTokens: 200,
- CacheCreation5mTokens: 80,
- CacheCreation1hTokens: 120,
- }
- result := calculateStatsCost(pricing, tokens, 1)
- require.NotNil(t, result)
- // 80*0.003 + 120*0.005 = 0.84
- require.InDelta(t, 0.84, *result, 1e-12)
- }
- func TestCalculateStatsCost_TokenBilling_WithImageOutput(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeToken,
- InputPrice: testPtrFloat64(0.001),
- OutputPrice: testPtrFloat64(0.002),
- ImageOutputPrice: testPtrFloat64(0.01),
- }
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- ImageOutputTokens: 10,
- }
- result := calculateStatsCost(pricing, tokens, 1)
- require.NotNil(t, result)
- // 100*0.001 + 50*0.002 + 10*0.01 = 0.1 + 0.1 + 0.1 = 0.3
- require.InDelta(t, 0.3, *result, 1e-12)
- }
- func TestCalculateStatsCost_TokenBilling_PartialPricesNil(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeToken,
- InputPrice: testPtrFloat64(0.001),
- // OutputPrice, CacheWritePrice, etc. are all nil → treated as 0
- }
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- CacheCreationTokens: 200,
- }
- result := calculateStatsCost(pricing, tokens, 1)
- require.NotNil(t, result)
- // Only input contributes: 100*0.001 = 0.1
- require.InDelta(t, 0.1, *result, 1e-12)
- }
- func TestCalculateStatsCost_TokenBilling_AllTokensZero(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeToken,
- InputPrice: testPtrFloat64(0.001),
- OutputPrice: testPtrFloat64(0.002),
- }
- tokens := UsageTokens{} // all zeros
- result := calculateStatsCost(pricing, tokens, 1)
- // totalCost == 0 → returns nil (does not override, falls back to default formula)
- require.Nil(t, result)
- }
- func TestCalculateStatsCost_PerRequestBilling(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModePerRequest,
- PerRequestPrice: testPtrFloat64(0.05),
- }
- tokens := UsageTokens{InputTokens: 999, OutputTokens: 999}
- result := calculateStatsCost(pricing, tokens, 3)
- require.NotNil(t, result)
- // 0.05 * 3 = 0.15
- require.InDelta(t, 0.15, *result, 1e-12)
- }
- func TestCalculateStatsCost_PerRequestBilling_PriceNil(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModePerRequest,
- // PerRequestPrice is nil
- }
- result := calculateStatsCost(pricing, UsageTokens{}, 1)
- require.Nil(t, result)
- }
- func TestCalculateStatsCost_PerRequestBilling_PriceZero(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModePerRequest,
- PerRequestPrice: testPtrFloat64(0),
- }
- result := calculateStatsCost(pricing, UsageTokens{}, 1)
- // price == 0 → condition *pricing.PerRequestPrice > 0 is false → returns nil
- require.Nil(t, result)
- }
- func TestCalculateStatsCost_ImageBilling(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeImage,
- PerRequestPrice: testPtrFloat64(0.10),
- }
- result := calculateStatsCost(pricing, UsageTokens{}, 2)
- require.NotNil(t, result)
- // 0.10 * 2 = 0.20
- require.InDelta(t, 0.20, *result, 1e-12)
- }
- func TestCalculateStatsCost_ImageBilling_PriceNil(t *testing.T) {
- pricing := &ChannelModelPricing{
- BillingMode: BillingModeImage,
- // PerRequestPrice is nil
- }
- result := calculateStatsCost(pricing, UsageTokens{}, 1)
- require.Nil(t, result)
- }
- func TestCalculateStatsCost_DefaultBillingMode_FallsToToken(t *testing.T) {
- // BillingMode is empty string (default) → falls into token billing
- pricing := &ChannelModelPricing{
- InputPrice: testPtrFloat64(0.001),
- OutputPrice: testPtrFloat64(0.002),
- }
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- }
- result := calculateStatsCost(pricing, tokens, 1)
- require.NotNil(t, result)
- require.InDelta(t, 0.2, *result, 1e-12)
- }
- // ---------------------------------------------------------------------------
- // tryCustomRules — 多规则顺序测试
- // ---------------------------------------------------------------------------
- func TestTryCustomRules_FirstMatchWins(t *testing.T) {
- channel := &Channel{
- AccountStatsPricingRules: []AccountStatsPricingRule{
- {
- GroupIDs: []int64{1},
- Pricing: []ChannelModelPricing{
- {ID: 100, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.01), OutputPrice: testPtrFloat64(0.02)},
- },
- },
- {
- GroupIDs: []int64{1},
- Pricing: []ChannelModelPricing{
- {ID: 200, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.99), OutputPrice: testPtrFloat64(0.99)},
- },
- },
- },
- }
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := tryCustomRules(channel, 999, 1, "", "claude-opus-4", tokens, 1)
- require.NotNil(t, result)
- // 应使用第一条规则的价格:100*0.01 + 50*0.02 = 2.0
- require.InDelta(t, 2.0, *result, 1e-12)
- }
- func TestTryCustomRules_SkipsNonMatchingRules(t *testing.T) {
- channel := &Channel{
- AccountStatsPricingRules: []AccountStatsPricingRule{
- {
- AccountIDs: []int64{888}, // 不匹配
- Pricing: []ChannelModelPricing{
- {ID: 100, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.99)},
- },
- },
- {
- GroupIDs: []int64{1}, // 匹配
- Pricing: []ChannelModelPricing{
- {ID: 200, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.05)},
- },
- },
- },
- }
- tokens := UsageTokens{InputTokens: 100}
- result := tryCustomRules(channel, 999, 1, "", "claude-opus-4", tokens, 1)
- require.NotNil(t, result)
- // 跳过规则1(账号不匹配),使用规则2:100*0.05 = 5.0
- require.InDelta(t, 5.0, *result, 1e-12)
- }
- func TestTryCustomRules_NoMatch_ReturnsNil(t *testing.T) {
- channel := &Channel{
- AccountStatsPricingRules: []AccountStatsPricingRule{
- {
- AccountIDs: []int64{888},
- Pricing: []ChannelModelPricing{
- {ID: 100, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.01)},
- },
- },
- },
- }
- tokens := UsageTokens{InputTokens: 100}
- result := tryCustomRules(channel, 999, 2, "", "claude-opus-4", tokens, 1)
- require.Nil(t, result) // 账号和分组都不匹配
- }
- func TestTryCustomRules_RuleMatchesButModelNot_ContinuesToNext(t *testing.T) {
- channel := &Channel{
- AccountStatsPricingRules: []AccountStatsPricingRule{
- {
- GroupIDs: []int64{1},
- Pricing: []ChannelModelPricing{
- {ID: 100, Models: []string{"gpt-4o"}, InputPrice: testPtrFloat64(0.01)}, // 模型不匹配
- },
- },
- {
- GroupIDs: []int64{1},
- Pricing: []ChannelModelPricing{
- {ID: 200, Models: []string{"claude-opus-4"}, InputPrice: testPtrFloat64(0.05)}, // 模型匹配
- },
- },
- },
- }
- tokens := UsageTokens{InputTokens: 100}
- result := tryCustomRules(channel, 999, 1, "", "claude-opus-4", tokens, 1)
- require.NotNil(t, result)
- require.InDelta(t, 5.0, *result, 1e-12) // 使用规则2
- }
- // ---------------------------------------------------------------------------
- // tryModelFilePricing
- // ---------------------------------------------------------------------------
- // newTestBillingServiceWithPrices creates a BillingService with pre-populated
- // fallback prices for testing. No config or pricing service is needed.
- // The key must match what getFallbackPricing resolves to for a given model name.
- // E.g., model "claude-sonnet-4" resolves to key "claude-sonnet-4".
- func newTestBillingServiceWithPrices(prices map[string]*ModelPricing) *BillingService {
- return &BillingService{
- fallbackPrices: prices,
- }
- }
- func TestTryModelFilePricing_Success(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-sonnet-4": {
- InputPricePerToken: 0.001,
- OutputPricePerToken: 0.002,
- },
- })
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
- require.NotNil(t, result)
- // 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
- require.InDelta(t, 0.2, *result, 1e-12)
- }
- func TestTryModelFilePricing_Fable51MaxEffortUsesTripleQuota(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-fable-5-1": {InputPricePerToken: 0.001},
- })
- tokens := UsageTokens{InputTokens: 100}
- standard := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, "xhigh")
- max := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, "max")
- require.NotNil(t, standard)
- require.NotNil(t, max)
- require.InDelta(t, *standard*3, *max, 1e-12)
- }
- func TestTryModelFilePricing_AppliesLongContextPricing(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "gpt-5.6-sol": {
- InputPricePerToken: 0.001,
- OutputPricePerToken: 0.002,
- CacheReadPricePerToken: 0.0001,
- LongContextInputThreshold: 100,
- LongContextInputMultiplier: 2,
- LongContextOutputMultiplier: 1.5,
- },
- })
- tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
- result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "", time.Time{})
- require.NotNil(t, result)
- // Input and cache-read use the 2x input tier; output uses the 1.5x tier.
- require.InDelta(t, 0.233, *result, 1e-12)
- }
- func TestTryModelFilePricing_AppliesServiceTierPricing(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "gpt-5.6-sol": {
- InputPricePerToken: 0.001,
- InputPricePerTokenPriority: 0.002,
- OutputPricePerToken: 0.002,
- OutputPricePerTokenPriority: 0.004,
- CacheCreationPricePerToken: 0.003,
- CacheCreationPricePerTokenPriority: 0.006,
- CacheReadPricePerToken: 0.0005,
- CacheReadPricePerTokenPriority: 0.001,
- },
- })
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- CacheCreationTokens: 20,
- CacheReadTokens: 10,
- }
- tests := []struct {
- name string
- serviceTier string
- want float64
- }{
- {name: "standard", serviceTier: "", want: 0.265},
- {name: "priority", serviceTier: "priority", want: 0.53},
- {name: "flex", serviceTier: "flex", want: 0.1325},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier, time.Time{})
- require.NotNil(t, result)
- require.InDelta(t, tt.want, *result, 1e-12)
- })
- }
- }
- func TestTryModelFilePricing_CombinesPriorityAndLongContextPricing(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "gpt-5.6-sol": {
- InputPricePerToken: 0.001,
- InputPricePerTokenPriority: 0.002,
- OutputPricePerToken: 0.002,
- OutputPricePerTokenPriority: 0.004,
- CacheCreationPricePerToken: 0.003,
- CacheCreationPricePerTokenPriority: 0.006,
- CacheReadPricePerToken: 0.0005,
- CacheReadPricePerTokenPriority: 0.001,
- LongContextInputThreshold: 100,
- LongContextInputMultiplier: 2,
- LongContextOutputMultiplier: 1.5,
- },
- })
- tokens := UsageTokens{
- InputTokens: 101,
- OutputTokens: 10,
- CacheCreationTokens: 5,
- CacheReadTokens: 5,
- }
- result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority", time.Time{})
- require.NotNil(t, result)
- // priority 单价先应用,再叠加长上下文输入 2x、输出 1.5x。
- require.InDelta(t, 0.534, *result, 1e-12)
- }
- func TestTryModelFilePricing_PricingNotFound(t *testing.T) {
- // "nonexistent-model" does not match any fallback pattern
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := tryModelFilePricing(bs, "nonexistent-model", tokens, "", time.Time{})
- require.Nil(t, result)
- }
- func TestTryModelFilePricing_NilFallback(t *testing.T) {
- // getFallbackPricing returns nil when key maps to nil
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-sonnet-4": nil,
- })
- tokens := UsageTokens{InputTokens: 100}
- result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
- require.Nil(t, result)
- }
- func TestTryModelFilePricing_ZeroCost(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-sonnet-4": {
- InputPricePerToken: 0.001,
- OutputPricePerToken: 0.002,
- },
- })
- tokens := UsageTokens{} // all zero tokens → cost = 0 → nil
- result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
- require.Nil(t, result)
- }
- func TestTryModelFilePricing_WithImageOutput(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-sonnet-4": {
- InputPricePerToken: 0.001,
- OutputPricePerToken: 0.002,
- ImageOutputPricePerToken: 0.01,
- },
- })
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- ImageOutputTokens: 10,
- }
- result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
- require.NotNil(t, result)
- // ImageOutputTokens 是 OutputTokens 的子集,先扣除再按图片单价计。
- // 100*0.001 + (50-10)*0.002 + 10*0.01 = 0.1 + 0.08 + 0.1 = 0.28
- require.InDelta(t, 0.28, *result, 1e-12)
- }
- func TestTryModelFilePricing_WithCacheTokens(t *testing.T) {
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-sonnet-4": {
- InputPricePerToken: 0.001,
- OutputPricePerToken: 0.002,
- CacheCreationPricePerToken: 0.003,
- CacheReadPricePerToken: 0.0005,
- },
- })
- tokens := UsageTokens{
- InputTokens: 100,
- OutputTokens: 50,
- CacheCreationTokens: 200,
- CacheReadTokens: 300,
- }
- result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
- require.NotNil(t, result)
- // 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
- // = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
- require.InDelta(t, 0.95, *result, 1e-12)
- }
- func TestTryModelFilePricing_DeepSeekPeakPricing(t *testing.T) {
- weekday := func(hour, minute int) time.Time {
- return time.Date(2026, time.August, 24, hour, minute, 0, 0, time.UTC)
- }
- for _, model := range []struct {
- name string
- input, output, cacheReadPrice float64
- }{
- {"deepseek-v4-flash", 2.2e-7, 6.6e-7, 7e-9},
- {"deepseek-v4-pro", 6.6e-7, 1.98e-6, 2.2e-8},
- } {
- for _, usage := range []struct {
- name string
- tokens UsageTokens
- }{
- {"input", UsageTokens{InputTokens: 1000}},
- {"output", UsageTokens{OutputTokens: 500}},
- {"cache_read", UsageTokens{CacheReadTokens: 1000}},
- {"mixed", UsageTokens{InputTokens: 1000, OutputTokens: 500, CacheReadTokens: 1000}},
- } {
- t.Run(model.name+"/"+usage.name, func(t *testing.T) {
- bs := newTestBillingService()
- tokens := usage.tokens
- baseCost := float64(tokens.InputTokens)*model.input +
- float64(tokens.OutputTokens)*model.output + float64(tokens.CacheReadTokens)*model.cacheReadPrice
- for _, slot := range []struct {
- name string
- at time.Time
- multiplier float64
- }{
- {"before_morning_peak", weekday(0, 59), 1},
- {"morning_peak_start", weekday(1, 0), 2},
- {"morning_peak_last_minute", weekday(3, 59), 2},
- {"morning_peak_end", weekday(4, 0), 1},
- {"afternoon_peak_start", weekday(6, 0), 2},
- {"afternoon_peak_last_minute", weekday(9, 59), 2},
- {"afternoon_peak_end", weekday(10, 0), 1},
- {"saturday", time.Date(2026, time.August, 22, 2, 0, 0, 0, time.UTC), 1},
- {"sunday", time.Date(2026, time.August, 23, 7, 0, 0, 0, time.UTC), 1},
- } {
- t.Run(slot.name, func(t *testing.T) {
- cost := tryModelFilePricing(bs, model.name, tokens, "", slot.at)
- require.NotNil(t, cost)
- require.InDelta(t, baseCost*slot.multiplier, *cost, 1e-12)
- })
- }
- })
- }
- }
- }
- func TestResolveAccountStatsCost_DeepSeekPricingPriority(t *testing.T) {
- peak := time.Date(2026, time.August, 24, 2, 0, 0, 0, time.UTC)
- for _, tt := range []struct {
- name string
- customRule bool
- applyPricing bool
- noChannel bool
- want float64
- }{
- {name: "catalog", want: 1000 * 2.2e-7 * 2},
- {name: "custom_rule", customRule: true, want: 1},
- {name: "custom_rule_before_customer_price", customRule: true, applyPricing: true, want: 1},
- {name: "customer_price", applyPricing: true, want: 0.75},
- {name: "no_channel", noChannel: true},
- } {
- t.Run(tt.name, func(t *testing.T) {
- channel := &Channel{
- ID: 1, Status: StatusActive, ApplyPricingToAccountStats: tt.applyPricing,
- ModelPricing: []ChannelModelPricing{{
- Models: []string{"deepseek-v4-flash"}, InputPrice: testPtrFloat64(0.02),
- }},
- }
- if tt.customRule {
- channel.AccountStatsPricingRules = []AccountStatsPricingRule{{
- AccountIDs: []int64{1},
- Pricing: []ChannelModelPricing{{
- Models: []string{"deepseek-v4-flash"}, InputPrice: testPtrFloat64(0.001),
- }},
- }}
- }
- cs := newTestChannelServiceForStats(t, channel, 10, PlatformDeepseek)
- groupID := int64(10)
- if tt.noChannel {
- groupID = 99
- }
- cost := resolveAccountStatsCost(context.Background(), cs, newTestBillingService(),
- 1, groupID, "deepseek-v4-flash", UsageTokens{InputTokens: 1000}, 1, 0.75, "", peak)
- if tt.noChannel {
- require.Nil(t, cost)
- return
- }
- require.NotNil(t, cost)
- require.InDelta(t, tt.want, *cost, 1e-12)
- })
- }
- }
- // ---------------------------------------------------------------------------
- // resolveAccountStatsCost — integration tests covering the 4-level priority chain
- // ---------------------------------------------------------------------------
- func TestResolveAccountStatsCost_NilChannelService(t *testing.T) {
- result := resolveAccountStatsCost(
- context.Background(),
- nil, // channelService is nil
- newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
- 1, 1, "claude-sonnet-4",
- UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
- )
- require.Nil(t, result)
- }
- func TestResolveAccountStatsCost_EmptyUpstreamModel(t *testing.T) {
- cs := newTestChannelServiceForStats(t, &Channel{
- ID: 1,
- Status: StatusActive,
- }, 1, "")
- result := resolveAccountStatsCost(
- context.Background(),
- cs,
- newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
- 1, 1, "", // empty upstream model
- UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
- )
- require.Nil(t, result)
- }
- func TestResolveAccountStatsCost_GetChannelForGroupReturnsNil(t *testing.T) {
- // Group 99 is NOT in the cache, so GetChannelForGroup returns nil
- cs := newTestChannelServiceForStats(t, &Channel{
- ID: 1,
- Status: StatusActive,
- }, 1, "")
- result := resolveAccountStatsCost(
- context.Background(),
- cs,
- newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
- 1, 99, "claude-sonnet-4", // groupID 99 has no channel
- UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
- )
- require.Nil(t, result)
- }
- func TestResolveAccountStatsCost_HitsCustomRule(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- AccountStatsPricingRules: []AccountStatsPricingRule{
- {
- GroupIDs: []int64{10},
- Pricing: []ChannelModelPricing{
- {
- ID: 100,
- Models: []string{"claude-sonnet-4"},
- InputPrice: testPtrFloat64(0.01),
- OutputPrice: testPtrFloat64(0.02),
- },
- },
- },
- },
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := resolveAccountStatsCost(
- context.Background(),
- cs, nil, // billingService not needed when custom rule hits
- 1, 10, "claude-sonnet-4",
- tokens, 1, 999.0, "priority", time.Time{}, // 自定义账号价格不叠加服务层级倍率
- )
- require.NotNil(t, result)
- // 100*0.01 + 50*0.02 = 1.0 + 1.0 = 2.0
- require.InDelta(t, 2.0, *result, 1e-12)
- }
- func TestResolveAccountStatsCost_ApplyPricingToAccountStats_UsesTotalCost(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: true,
- // No custom rules
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := resolveAccountStatsCost(
- context.Background(),
- cs, nil,
- 1, 10, "claude-sonnet-4",
- tokens, 1, 0.75, "priority", time.Time{}, // 已完成用户计费,不再重复应用服务层级倍率
- )
- require.NotNil(t, result)
- require.InDelta(t, 0.75, *result, 1e-12)
- }
- func TestResolveAccountStatsCost_ApplyPricingToAccountStats_ZeroTotalCost_ReturnsNil(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: true,
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- result := resolveAccountStatsCost(
- context.Background(),
- cs, nil,
- 1, 10, "claude-sonnet-4",
- UsageTokens{}, 1, 0.0, "", time.Time{}, // totalCost = 0
- )
- require.Nil(t, result)
- }
- func TestResolveAccountStatsCost_FallsBackToLiteLLM(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: false, // not enabled
- // No custom rules
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-sonnet-4": {
- InputPricePerToken: 0.001,
- OutputPricePerToken: 0.002,
- },
- })
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := resolveAccountStatsCost(
- context.Background(),
- cs, bs,
- 1, 10, "claude-sonnet-4",
- tokens, 1, 999.0, "", time.Time{}, // totalCost ignored
- )
- require.NotNil(t, result)
- // 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
- require.InDelta(t, 0.2, *result, 1e-12)
- }
- func TestResolveAccountStatsCost_FallbackHonorsAnthropicFast(t *testing.T) {
- channel := &Channel{ID: 1, Status: StatusActive}
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "claude-opus-5": {
- InputPricePerToken: 5e-6,
- OutputPricePerToken: 25e-6,
- },
- })
- result := resolveAccountStatsCost(
- context.Background(), cs, bs,
- 1, 10, "claude-opus-5",
- UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000},
- 1, 0, "fast", time.Time{},
- )
- require.NotNil(t, result)
- require.InDelta(t, 60, *result, 1e-12)
- }
- func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: false,
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "antigravity")
- bs := NewBillingService(&config.Config{}, nil)
- result := resolveAccountStatsCost(
- context.Background(),
- cs, bs,
- 1, 10, "gemini-3.6-flash-low",
- UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{},
- )
- require.NotNil(t, result)
- require.InDelta(t, 9.15, *result, 1e-12)
- }
- func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: false,
- // No custom rules
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- // BillingService with no pricing for the model
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
- tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
- result := resolveAccountStatsCost(
- context.Background(),
- cs, bs,
- 1, 10, "totally-unknown-model",
- tokens, 1, 0.0, "", time.Time{},
- )
- require.Nil(t, result)
- }
- func TestResolveAccountStatsCost_NilBillingService_SkipsLiteLLM(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: false,
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- result := resolveAccountStatsCost(
- context.Background(),
- cs, nil, // billingService is nil
- 1, 10, "claude-sonnet-4",
- UsageTokens{InputTokens: 100}, 1, 0.0, "", time.Time{},
- )
- require.Nil(t, result)
- }
- func TestResolveAccountStatsCost_CustomRulePriorityOverApplyPricing(t *testing.T) {
- // Both custom rule and ApplyPricingToAccountStats are configured;
- // custom rule should take precedence.
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: true,
- AccountStatsPricingRules: []AccountStatsPricingRule{
- {
- GroupIDs: []int64{10},
- Pricing: []ChannelModelPricing{
- {
- ID: 100,
- Models: []string{"claude-sonnet-4"},
- InputPrice: testPtrFloat64(0.05),
- },
- },
- },
- },
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "anthropic")
- tokens := UsageTokens{InputTokens: 100}
- result := resolveAccountStatsCost(
- context.Background(),
- cs, nil,
- 1, 10, "claude-sonnet-4",
- tokens, 1, 99.0, "", time.Time{}, // totalCost = 99.0 (would be used if ApplyPricing wins)
- )
- require.NotNil(t, result)
- // Custom rule: 100*0.05 = 5.0 (NOT 99.0 from totalCost)
- require.InDelta(t, 5.0, *result, 1e-12)
- }
- func TestApplyAccountStatsCost_UsesUsageLogServiceTier(t *testing.T) {
- channel := &Channel{
- ID: 1,
- Status: StatusActive,
- ApplyPricingToAccountStats: false,
- }
- cs := newTestChannelServiceForStats(t, channel, 10, "openai")
- bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
- "gpt-5.6-sol": {
- InputPricePerToken: 0.001,
- InputPricePerTokenPriority: 0.002,
- OutputPricePerToken: 0.002,
- OutputPricePerTokenPriority: 0.004,
- },
- })
- serviceTier := "priority"
- usageLog := &UsageLog{ServiceTier: &serviceTier}
- applyAccountStatsCost(
- context.Background(), usageLog, cs, bs,
- 1, 10, "gpt-5.6-sol", "gpt-5.6-sol",
- UsageTokens{InputTokens: 100, OutputTokens: 50}, 999, time.Time{},
- )
- require.NotNil(t, usageLog.AccountStatsCost)
- require.InDelta(t, 0.4, *usageLog.AccountStatsCost, 1e-12)
- }
- // ---------------------------------------------------------------------------
- // helpers for resolveAccountStatsCost tests
- // ---------------------------------------------------------------------------
- // newTestChannelServiceForStats creates a ChannelService with a single channel
- // mapped to the given groupID, suitable for resolveAccountStatsCost tests.
- func newTestChannelServiceForStats(t *testing.T, channel *Channel, groupID int64, platform string) *ChannelService {
- t.Helper()
- cache := newEmptyChannelCache()
- cache.channelByGroupID[groupID] = channel
- cache.groupPlatform[groupID] = platform
- cs := &ChannelService{}
- cache.loadedAt = time.Now()
- cs.cache.Store(cache)
- return cs
- }
|