package helper import ( "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/types" "github.com/stretchr/testify/require" ) func TestPricingLookup_MatchSpecificity(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackReject) resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "720P", "duration": 10.0}} result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, taskErr) require.Equal(t, types.PricingModeMatched, result.Mode) require.Equal(t, 0.08, result.PriceUSD) require.Equal(t, "720P", result.Snapshot["resolution"]) require.Equal(t, 10.0, result.Snapshot["duration"]) } func TestPricingLookup_StringDimensionMatchesCaseInsensitive(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackReject) resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "720p", "duration": 10.0}} result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, taskErr) require.Equal(t, types.PricingModeMatched, result.Mode) require.Equal(t, 0.08, result.PriceUSD) require.Equal(t, "720p", result.Snapshot["resolution"]) } func TestPricingLookup_WildcardFallbackMatch(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackReject) resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "480P", "duration": 10.0}} result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, taskErr) require.Equal(t, 0.05, result.PriceUSD) } func TestPricingLookup_RejectFallback(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackReject) resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "1080P", "duration": 15.0}} result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, result) require.NotNil(t, taskErr) require.Equal(t, "pricing_no_match", taskErr.Code) } func TestPricingLookup_MaxFallback(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackMax) resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "1080P", "duration": 15.0}} result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, taskErr) require.Equal(t, types.PricingModeFallbackMax, result.Mode) require.Equal(t, 0.1, result.PriceUSD) require.Equal(t, types.PricingFallbackMax, result.Snapshot["fallback_strategy"]) } func TestPricingLookup_DefaultFallback(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackDefault) cfg.Fallback.DefaultPrice = 0.2 resolved := &ResolvedDimensionValues{Values: map[string]any{"resolution": "1080P", "duration": 15.0}} result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, taskErr) require.Equal(t, types.PricingModeFallbackDefault, result.Mode) require.Equal(t, 0.2, result.PriceUSD) } func TestPricingLookup_MissingDimensionSnapshotMarshalsNull(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackDefault) cfg.Fallback.DefaultPrice = 0.2 resolved := &ResolvedDimensionValues{ Values: map[string]any{"resolution": "1080P", "duration": nil}, Missing: map[string]bool{"duration": true}, } result, taskErr := lookupPricingConfig(cfg, resolved) require.Nil(t, taskErr) body, err := common.Marshal(result.Snapshot) require.NoError(t, err) require.Contains(t, string(body), `"duration":null`) } func TestPricingLookup_DecisionQuota(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackReject) result := &pricingLookupResult{Mode: types.PricingModeMatched, PriceUSD: 0.1, Snapshot: map[string]any{"price": 0.1}} decision := buildMatchedDecision(cfg, result, types.GroupRatioInfo{GroupRatio: 0.5}) require.Equal(t, types.BillingModeMatrix, decision.BillingMode) require.Equal(t, types.BillingUnitPerCall, decision.BillingUnit) require.Equal(t, int(0.1*common.QuotaPerUnit*0.5), decision.Quota) require.True(t, decision.PerCallBilling) } func TestPricingLookup_UsageDecisionMinimumPreconsume(t *testing.T) { cfg := testPricingLookupConfig(types.PricingFallbackReject) cfg.BillingUnit = types.BillingUnitPer1MTokens cfg.PreconsumeStrategy = types.PreconsumeStrategyMinimum result := &pricingLookupResult{Mode: types.PricingModeMatched, PriceUSD: 2.0, Snapshot: map[string]any{"price": 2.0}} decision := buildMatchedDecision(cfg, result, types.GroupRatioInfo{GroupRatio: 1}) require.Equal(t, 2.0, decision.PriceUSD) require.Equal(t, 2.0, decision.TokenUnitPriceUSD) require.Equal(t, minimumUsagePreconsumeQuota(2.0, 1), decision.Quota) require.False(t, decision.PerCallBilling) } func TestPricingLookup_RemixUsageDecisionKeepsDisplayPrice(t *testing.T) { decision, taskErr := buildRemixMatrixDecision(&types.OriginPricingSnapshot{ PriceUSD: 0, TokenUnitPriceUSD: 46, BillingUnit: types.BillingUnitPer1MTokens, Snapshot: map[string]any{"price": 46}, PerCallBilling: false, }, types.GroupRatioInfo{GroupRatio: 1}) require.Nil(t, taskErr) require.Equal(t, 46.0, decision.PriceUSD) require.Equal(t, 46.0, decision.TokenUnitPriceUSD) require.Equal(t, minimumUsagePreconsumeQuota(46, 1), decision.Quota) require.False(t, decision.PerCallBilling) } func testPricingLookupConfig(fallback string) *types.PricingConfig { return &types.PricingConfig{ SchemaVersion: types.PricingSchemaVersion, Scope: types.PricingScopeModel, BillingUnit: types.BillingUnitPerCall, PreconsumeStrategy: types.PreconsumeStrategyExact, Dimensions: []types.PricingDimension{ {Key: "resolution", Source: "request.resolution", Type: "string"}, {Key: "duration", Source: "request.duration", Type: "number"}, }, Table: []types.PricingRow{ {"resolution": "720P", "duration": 10.0, "price": 0.08, "source": types.PricingRowSourceManual}, {"resolution": "720P", "duration": "*", "price": 0.1, "source": types.PricingRowSourceManual}, {"resolution": "*", "duration": 10.0, "price": 0.05, "source": types.PricingRowSourceManual}, }, Fallback: types.PricingFallback{Strategy: fallback}, } }