new-api/relay/helper/pricing_lookup_test.go

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_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.TokenUnitPriceUSD)
	require.Equal(t, minimumUsagePreconsumeQuota(2.0, 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},
	}
}