| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414 |
- //go:build unit
- package service
- import (
- "context"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/stretchr/testify/require"
- )
- // ---------------------------------------------------------------------------
- // CalculateCostUnified
- // ---------------------------------------------------------------------------
- func TestCalculateCostUnified_NilResolver_FallsBackToOldPath(t *testing.T) {
- svc := newTestBillingService()
- tokens := UsageTokens{InputTokens: 1000, OutputTokens: 500}
- input := CostInput{
- Model: "claude-sonnet-4",
- Tokens: tokens,
- RateMultiplier: 1.0,
- Resolver: nil, // no resolver
- }
- cost, err := svc.CalculateCostUnified(input)
- require.NoError(t, err)
- // Should match the old-path result exactly
- expected, err := svc.calculateCostInternal("claude-sonnet-4", tokens, 1.0, "", nil)
- require.NoError(t, err)
- require.InDelta(t, expected.TotalCost, cost.TotalCost, 1e-10)
- require.InDelta(t, expected.ActualCost, cost.ActualCost, 1e-10)
- // BillingMode is NOT set by old path through CalculateCostUnified (resolver == nil)
- require.Empty(t, cost.BillingMode)
- }
- func TestCalculateCostUnified_TokenMode(t *testing.T) {
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(nil, bs)
- tokens := UsageTokens{InputTokens: 1000, OutputTokens: 500}
- input := CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- Tokens: tokens,
- RateMultiplier: 1.5,
- Resolver: resolver,
- }
- cost, err := bs.CalculateCostUnified(input)
- require.NoError(t, err)
- require.NotNil(t, cost)
- // Verify token billing: Input: 1000*3e-6=0.003, Output: 500*15e-6=0.0075
- expectedTotal := 1000*3e-6 + 500*15e-6
- require.InDelta(t, expectedTotal, cost.TotalCost, 1e-10)
- require.InDelta(t, expectedTotal*1.5, cost.ActualCost, 1e-10)
- require.Equal(t, string(BillingModeToken), cost.BillingMode)
- }
- func TestCalculateCostUnified_TokenModeAppliesRateMultiplierToImageTokens(t *testing.T) {
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(nil, bs)
- tokens := UsageTokens{InputTokens: 1000, OutputTokens: 600, ImageOutputTokens: 100}
- cost, err := bs.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- Tokens: tokens,
- RateMultiplier: 3.0,
- Resolver: resolver,
- })
- require.NoError(t, err)
- textInput := 1000 * 3e-6
- textOutput := 500 * 15e-6
- imageOutput := 100 * 15e-6
- require.InDelta(t, textInput+textOutput+imageOutput, cost.TotalCost, 1e-10)
- require.InDelta(t, (textInput+textOutput+imageOutput)*3.0, cost.ActualCost, 1e-10)
- require.InDelta(t, imageOutput, cost.ImageOutputCost, 1e-10)
- }
- func TestCalculateCostUnified_PerRequestMode(t *testing.T) {
- // Set up a ChannelService with a per-request pricing channel
- cs := newTestChannelServiceWithCache(t, &channelCache{
- pricingByGroupModel: map[channelModelKey]*ChannelModelPricing{
- {groupID: 1, model: "claude-sonnet-4"}: {
- BillingMode: BillingModePerRequest,
- PerRequestPrice: testPtrFloat64(0.05),
- },
- },
- channelByGroupID: map[int64]*Channel{
- 1: {ID: 1, Status: StatusActive},
- },
- groupPlatform: map[int64]string{1: ""},
- wildcardByGroupPlatform: map[channelGroupPlatformKey][]*wildcardPricingEntry{},
- mappingByGroupModel: map[channelModelKey]string{},
- wildcardMappingByGP: map[channelGroupPlatformKey][]*wildcardMappingEntry{},
- byID: map[int64]*Channel{},
- })
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(cs, bs)
- groupID := int64(1)
- input := CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- GroupID: &groupID,
- Tokens: UsageTokens{InputTokens: 100, OutputTokens: 50},
- RequestCount: 3,
- RateMultiplier: 2.0,
- Resolver: resolver,
- }
- cost, err := bs.CalculateCostUnified(input)
- require.NoError(t, err)
- require.NotNil(t, cost)
- // 3 requests * $0.05 = $0.15
- require.InDelta(t, 0.15, cost.TotalCost, 1e-10)
- // ActualCost = 0.15 * 2.0 = 0.30
- require.InDelta(t, 0.30, cost.ActualCost, 1e-10)
- require.Equal(t, string(BillingModePerRequest), cost.BillingMode)
- }
- func TestCalculateCostUnified_ImageMode(t *testing.T) {
- cs := newTestChannelServiceWithCache(t, &channelCache{
- pricingByGroupModel: map[channelModelKey]*ChannelModelPricing{
- {groupID: 2, model: "gemini-image"}: {
- BillingMode: BillingModeImage,
- PerRequestPrice: testPtrFloat64(0.10),
- },
- },
- channelByGroupID: map[int64]*Channel{
- 2: {ID: 2, Status: StatusActive},
- },
- groupPlatform: map[int64]string{2: ""},
- wildcardByGroupPlatform: map[channelGroupPlatformKey][]*wildcardPricingEntry{},
- mappingByGroupModel: map[channelModelKey]string{},
- wildcardMappingByGP: map[channelGroupPlatformKey][]*wildcardMappingEntry{},
- byID: map[int64]*Channel{},
- })
- bs := &BillingService{
- cfg: &config.Config{},
- fallbackPrices: map[string]*ModelPricing{},
- }
- resolver := NewModelPricingResolver(cs, bs)
- groupID := int64(2)
- input := CostInput{
- Ctx: context.Background(),
- Model: "gemini-image",
- GroupID: &groupID,
- Tokens: UsageTokens{},
- RequestCount: 2,
- RateMultiplier: 1.0,
- Resolver: resolver,
- }
- cost, err := bs.CalculateCostUnified(input)
- require.NoError(t, err)
- require.NotNil(t, cost)
- // 2 * $0.10 = $0.20
- require.InDelta(t, 0.20, cost.TotalCost, 1e-10)
- require.InDelta(t, 0.20, cost.ActualCost, 1e-10)
- require.Equal(t, string(BillingModeImage), cost.BillingMode)
- }
- func channelTimeResolvedForTest(base *ModelPricing, intervals []PricingInterval) *ResolvedPricing {
- return &ResolvedPricing{
- Mode: BillingModeToken,
- BasePricing: base,
- Intervals: intervals,
- Source: PricingSourceChannel,
- channelPricing: &ChannelModelPricing{
- BillingMode: BillingModeToken,
- TimePricing: &ChannelTimePricing{
- Timezone: "Asia/Shanghai",
- Periods: []ChannelTimePricingPeriod{{
- StartTime: "09:00",
- EndTime: "12:00",
- Multiplier: 2,
- }},
- },
- },
- longContextPricingEnabled: true,
- }
- }
- func TestCalculateCostUnified_ChannelTimePricingScalesBaseAndActualCost(t *testing.T) {
- billing := NewBillingService(&config.Config{}, nil)
- resolved := channelTimeResolvedForTest(&ModelPricing{InputPricePerToken: 0.001}, nil)
- cost, err := billing.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "model",
- Tokens: UsageTokens{InputTokens: 1000},
- RateMultiplier: 0.8,
- Resolver: &ModelPricingResolver{},
- Resolved: resolved,
- PricingAt: time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
- })
- require.NoError(t, err)
- require.InDelta(t, 2.0, cost.InputCost, 1e-12)
- require.InDelta(t, 2.0, cost.TotalCost, 1e-12)
- require.InDelta(t, 1.6, cost.ActualCost, 1e-12)
- }
- func TestCalculateCostUnified_ChannelTimePricingScalesMatchingInterval(t *testing.T) {
- intervalInputPrice := 0.003
- resolved := channelTimeResolvedForTest(
- &ModelPricing{InputPricePerToken: 0.001},
- []PricingInterval{{MinTokens: 0, InputPrice: &intervalInputPrice}},
- )
- billing := NewBillingService(&config.Config{}, nil)
- cost, err := billing.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "model",
- Tokens: UsageTokens{InputTokens: 1000},
- Resolver: &ModelPricingResolver{},
- Resolved: resolved,
- PricingAt: time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
- })
- require.NoError(t, err)
- require.InDelta(t, 6.0, cost.InputCost, 1e-12)
- require.InDelta(t, 6.0, cost.TotalCost, 1e-12)
- }
- func TestCalculateCostUnified_ChannelTimePricingScalesBaseOnUnmatchedInterval(t *testing.T) {
- intervalInputPrice := 0.003
- resolved := channelTimeResolvedForTest(
- &ModelPricing{InputPricePerToken: 0.001},
- []PricingInterval{{MinTokens: 2000, InputPrice: &intervalInputPrice}},
- )
- billing := NewBillingService(&config.Config{}, nil)
- cost, err := billing.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "model",
- Tokens: UsageTokens{InputTokens: 1000},
- Resolver: &ModelPricingResolver{},
- Resolved: resolved,
- PricingAt: time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
- })
- require.NoError(t, err)
- require.InDelta(t, 2.0, cost.InputCost, 1e-12)
- require.InDelta(t, 2.0, cost.TotalCost, 1e-12)
- }
- func TestCalculateCostUnified_ChannelTimePricingDoesNotApplyToGroupPricing(t *testing.T) {
- resolved := channelTimeResolvedForTest(&ModelPricing{InputPricePerToken: 0.001}, nil)
- resolved.Source = PricingSourceGroup
- billing := NewBillingService(&config.Config{}, nil)
- cost, err := billing.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "model",
- Tokens: UsageTokens{InputTokens: 1000},
- Resolver: &ModelPricingResolver{},
- Resolved: resolved,
- PricingAt: time.Date(2026, 8, 17, 1, 0, 0, 0, time.UTC),
- })
- require.NoError(t, err)
- require.InDelta(t, 1.0, cost.TotalCost, 1e-12)
- }
- func TestCalculateCostUnified_ChannelTimePricingDoesNotApplyOutsideMatchingTime(t *testing.T) {
- resolved := channelTimeResolvedForTest(&ModelPricing{InputPricePerToken: 0.001}, nil)
- billing := NewBillingService(&config.Config{}, nil)
- for _, pricingAt := range []time.Time{
- time.Time{},
- time.Date(2026, 8, 17, 5, 0, 0, 0, time.UTC),
- } {
- cost, err := billing.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "model",
- Tokens: UsageTokens{InputTokens: 1000},
- Resolver: &ModelPricingResolver{},
- Resolved: resolved,
- PricingAt: pricingAt,
- })
- require.NoError(t, err)
- require.InDelta(t, 1.0, cost.TotalCost, 1e-12)
- }
- }
- func TestApplyCostBreakdownMultiplierScalesAllMonetaryFields(t *testing.T) {
- cost := &CostBreakdown{
- InputCost: 1,
- ImageInputCost: 2,
- OutputCost: 3,
- ImageOutputCost: 4,
- CacheCreationCost: 5,
- CacheReadCost: 6,
- TotalCost: 21,
- ActualCost: 42,
- BillingMode: string(BillingModeToken),
- LongContextBillingApplied: true,
- }
- applyCostBreakdownMultiplier(cost, 1.5)
- require.InDelta(t, 1.5, cost.InputCost, 1e-12)
- require.InDelta(t, 3.0, cost.ImageInputCost, 1e-12)
- require.InDelta(t, 4.5, cost.OutputCost, 1e-12)
- require.InDelta(t, 6.0, cost.ImageOutputCost, 1e-12)
- require.InDelta(t, 7.5, cost.CacheCreationCost, 1e-12)
- require.InDelta(t, 9.0, cost.CacheReadCost, 1e-12)
- require.InDelta(t, 31.5, cost.TotalCost, 1e-12)
- require.InDelta(t, 63.0, cost.ActualCost, 1e-12)
- require.Equal(t, string(BillingModeToken), cost.BillingMode)
- require.True(t, cost.LongContextBillingApplied)
- }
- // TestCalculateCostUnified_RateMultiplierZeroProducesZero 锁定新行为:
- // 保存时强制 > 0;若 0 仍泄漏到计费层,按 0 计费(而非历史上的 1.0)。
- func TestCalculateCostUnified_RateMultiplierZeroProducesZero(t *testing.T) {
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(nil, bs)
- tokens := UsageTokens{InputTokens: 1000, OutputTokens: 500}
- cost, err := bs.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- Tokens: tokens,
- RateMultiplier: 0,
- Resolver: resolver,
- })
- require.NoError(t, err)
- require.Greater(t, cost.TotalCost, 0.0)
- require.InDelta(t, 0.0, cost.ActualCost, 1e-10)
- }
- // TestCalculateCostUnified_NegativeRateMultiplierClampedToZero 锁定新行为:
- // 负数倍率按 0 计费,避免历史的 <=0 → 1.0 把配置异常静默按标准价扣费。
- func TestCalculateCostUnified_NegativeRateMultiplierClampedToZero(t *testing.T) {
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(nil, bs)
- tokens := UsageTokens{InputTokens: 1000}
- cost, err := bs.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- Tokens: tokens,
- RateMultiplier: -5.0,
- Resolver: resolver,
- })
- require.NoError(t, err)
- require.Greater(t, cost.TotalCost, 0.0)
- require.InDelta(t, 0.0, cost.ActualCost, 1e-10)
- }
- func TestCalculateCostUnified_BillingModeFieldFilled(t *testing.T) {
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(nil, bs)
- cost, err := bs.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- Tokens: UsageTokens{InputTokens: 100},
- RateMultiplier: 1.0,
- Resolver: resolver,
- })
- require.NoError(t, err)
- require.Equal(t, "token", cost.BillingMode)
- }
- func TestCalculateCostUnified_UsesPreResolvedPricing(t *testing.T) {
- bs := newTestBillingService()
- resolver := NewModelPricingResolver(nil, bs)
- // Pre-resolve with per_request mode to verify it's used instead of re-resolving
- preResolved := &ResolvedPricing{
- Mode: BillingModePerRequest,
- DefaultPerRequestPrice: 0.07,
- }
- cost, err := bs.CalculateCostUnified(CostInput{
- Ctx: context.Background(),
- Model: "claude-sonnet-4",
- Tokens: UsageTokens{InputTokens: 100},
- RequestCount: 2,
- RateMultiplier: 1.0,
- Resolver: resolver,
- Resolved: preResolved,
- })
- require.NoError(t, err)
- require.NotNil(t, cost)
- // 2 * $0.07 = $0.14
- require.InDelta(t, 0.14, cost.TotalCost, 1e-10)
- require.Equal(t, string(BillingModePerRequest), cost.BillingMode)
- }
- // ---------------------------------------------------------------------------
- // helpers
- // ---------------------------------------------------------------------------
- // newTestChannelServiceWithCache creates a ChannelService with a pre-populated
- // cache snapshot, bypassing the repository layer entirely.
- func newTestChannelServiceWithCache(t *testing.T, cache *channelCache) *ChannelService {
- t.Helper()
- cs := &ChannelService{}
- cache.loadedAt = time.Now()
- cs.cache.Store(cache)
- return cs
- }
|