new-api/setting/ratio_setting/model_pricing_test.go

package ratio_setting

import (
	"testing"

	"github.com/QuantumNous/new-api/types"
	"github.com/stretchr/testify/require"
)

func validPricingConfig() *types.PricingConfig {
	return &types.PricingConfig{
		SchemaVersion: types.PricingSchemaVersion,
		Scope:         types.PricingScopeModel,
		Dimensions: []types.PricingDimension{
			{Key: "resolution", Source: "request.resolution", Type: "string", Default: nil},
			{Key: "duration", Source: "request.duration", Type: "number", Default: nil},
		},
		Table: []types.PricingRow{
			{"resolution": "720P", "duration": float64(5), "price": float64(0.10), "source": types.PricingRowSourceManual},
			{"resolution": "720P", "duration": float64(10), "price": float64(0.20), "source": types.PricingRowSourceGenerated},
		},
		Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject},
	}
}

func TestValidatePricingConfigValidDefaults(t *testing.T) {
	cfg := validPricingConfig()
	require.NoError(t, ValidatePricingConfig("hailuo-test", cfg))
	require.Equal(t, types.BillingUnitPerCall, cfg.BillingUnit)
	require.Equal(t, types.PreconsumeStrategyExact, cfg.PreconsumeStrategy)
}

func TestValidatePricingConfigSchemaRules(t *testing.T) {
	tests := []struct {
		name    string
		mutate  func(*types.PricingConfig)
		wantErr string
	}{
		{"schema_version", func(c *types.PricingConfig) { c.SchemaVersion = 2 }, "schema_version"},
		{"scope", func(c *types.PricingConfig) { c.Scope = "channel" }, "scope"},
		{"dimensions", func(c *types.PricingConfig) { c.Dimensions = nil }, "dimensions"},
		{"table", func(c *types.PricingConfig) { c.Table = nil }, "table"},
		{"fallback", func(c *types.PricingConfig) { c.Fallback.Strategy = "nearest" }, "fallback.strategy"},
		{"default_price", func(c *types.PricingConfig) {
			c.Fallback = types.PricingFallback{Strategy: types.PricingFallbackDefault, DefaultPrice: -1}
		}, "default_price"},
		{"per_call_strategy", func(c *types.PricingConfig) {
			c.BillingUnit = types.BillingUnitPerCall
			c.PreconsumeStrategy = types.PreconsumeStrategyMinimum
		}, "per_call"},
		{"usage_strategy", func(c *types.PricingConfig) {
			c.BillingUnit = types.BillingUnitPer1MTokens
			c.PreconsumeStrategy = types.PreconsumeStrategyExact
		}, "per_1m_tokens"},
		{"usage_allowlist", func(c *types.PricingConfig) {
			c.BillingUnit = types.BillingUnitPer1MTokens
			c.PreconsumeStrategy = types.PreconsumeStrategyMinimum
		}, "does not support"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cfg := validPricingConfig()
			tt.mutate(cfg)
			err := ValidatePricingConfig("hailuo-test", cfg)
			require.Error(t, err)
			require.Contains(t, err.Error(), tt.wantErr)
		})
	}
}

func TestValidatePricingConfigUsageAllowlist(t *testing.T) {
	cfg := validPricingConfig()
	cfg.BillingUnit = types.BillingUnitPer1MTokens
	cfg.PreconsumeStrategy = types.PreconsumeStrategyMinimum
	require.NoError(t, ValidatePricingConfig("seedance-2", cfg))
}

func TestValidatePricingConfigUsageAllowlistKlingAipingModels(t *testing.T) {
	models := []string{
		"kling-v1",
		"kling-v1-6",
		"kling-v2-6",
		"kling-v3",
		"kling-video-o1",
		"kling-v3-omni",
	}

	for _, model := range models {
		t.Run(model, func(t *testing.T) {
			cfg := validPricingConfig()
			cfg.BillingUnit = types.BillingUnitPer1MTokens
			cfg.PreconsumeStrategy = types.PreconsumeStrategyMinimum
			require.NoError(t, ValidatePricingConfig(model, cfg))
		})
	}
}

func TestValidatePricingConfigAllowsDerivedVideoInput(t *testing.T) {
	cfg := &types.PricingConfig{
		SchemaVersion:      types.PricingSchemaVersion,
		Scope:              types.PricingScopeModel,
		BillingUnit:        types.BillingUnitPer1MTokens,
		PreconsumeStrategy: types.PreconsumeStrategyMinimum,
		Dimensions: []types.PricingDimension{
			{Key: "video_input", Source: "derived.video_input", Type: "boolean", Default: nil},
		},
		Table: []types.PricingRow{
			{"video_input": false, "price": float64(46), "source": types.PricingRowSourceManual},
			{"video_input": true, "price": float64(28), "source": types.PricingRowSourceManual},
		},
		Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject},
	}

	require.NoError(t, ValidatePricingConfig("seedance-2", cfg))
}

func TestValidatePricingConfigDimensionRules(t *testing.T) {
	tests := []struct {
		name    string
		dim     types.PricingDimension
		wantErr string
	}{
		{"bad_key", types.PricingDimension{Key: "Bad", Source: "request.x", Type: "string", Default: nil}, "invalid"},
		{"reserved", types.PricingDimension{Key: "price", Source: "request.x", Type: "string", Default: nil}, "reserved"},
		{"bad_type", types.PricingDimension{Key: "quality", Source: "request.x", Type: "object", Default: nil}, "type"},
		{"bad_source", types.PricingDimension{Key: "quality", Source: "response.usage.total_tokens", Type: "number", Default: nil}, "source"},
		{"optional_default_missing", types.PricingDimension{Key: "quality", Source: "request.quality", Type: "string", Optional: true, Default: nil}, "default required"},
		{"optional_default_type", types.PricingDimension{Key: "quality", Source: "request.quality", Type: "number", Optional: true, Default: "fast"}, "default type"},
		{"required_default", types.PricingDimension{Key: "quality", Source: "request.quality", Type: "string", Default: "fast"}, "default must be null"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cfg := validPricingConfig()
			cfg.Dimensions = []types.PricingDimension{tt.dim}
			cfg.Table = []types.PricingRow{{tt.dim.Key: "x", "price": float64(1), "source": types.PricingRowSourceManual}}
			if tt.dim.Type == "number" {
				cfg.Table[0][tt.dim.Key] = float64(1)
			}
			err := ValidatePricingConfig("hailuo-test", cfg)
			require.Error(t, err)
			require.Contains(t, err.Error(), tt.wantErr)
		})
	}
}

func TestValidatePricingConfigRowRules(t *testing.T) {
	tests := []struct {
		name    string
		mutate  func(*types.PricingConfig)
		wantErr string
	}{
		{"missing_dim", func(c *types.PricingConfig) { delete(c.Table[0], "duration") }, "missing dimension"},
		{"extra_field", func(c *types.PricingConfig) { c.Table[0]["price_formula"] = "x" }, "unexpected field"},
		{"bad_dim_type", func(c *types.PricingConfig) { c.Table[0]["duration"] = "five" }, "value type"},
		{"bad_price", func(c *types.PricingConfig) { c.Table[0]["price"] = -1.0 }, "price"},
		{"bad_source", func(c *types.PricingConfig) { c.Table[0]["source"] = "api" }, "source"},
		{"duplicate", func(c *types.PricingConfig) {
			c.Table = append(c.Table, types.PricingRow(types.CloneMapAny(c.Table[0])))
		}, "duplicate"},
		{"ambiguous", func(c *types.PricingConfig) {
			c.Dimensions = []types.PricingDimension{
				{Key: "res", Source: "request.res", Type: "string", Default: nil},
				{Key: "dur", Source: "request.dur", Type: "number", Default: nil},
			}
			c.Table = []types.PricingRow{
				{"res": "*", "dur": float64(5), "price": float64(0.10), "source": types.PricingRowSourceManual},
				{"res": "720P", "dur": "*", "price": float64(0.12), "source": types.PricingRowSourceManual},
			}
		}, "ambiguous"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cfg := validPricingConfig()
			tt.mutate(cfg)
			err := ValidatePricingConfig("hailuo-test", cfg)
			require.Error(t, err)
			require.Contains(t, err.Error(), tt.wantErr)
		})
	}
}

func TestModelPricingRulesRoundTripAndClone(t *testing.T) {
	require.NoError(t, UpdateModelPricingRulesByJSONString(`{}`))
	cfg := validPricingConfig()
	require.NoError(t, SetPricingConfig("hailuo-test", cfg))

	cfg.Table[0]["price"] = float64(999)
	got := GetPricingConfig("hailuo-test")
	require.NotNil(t, got)
	require.Equal(t, float64(0.10), got.Table[0]["price"])

	jsonStr := ModelPricingRules2JSONString()
	require.Contains(t, jsonStr, "hailuo-test")
	require.NoError(t, UpdateModelPricingRulesByJSONString(jsonStr))
	require.NotNil(t, GetPricingConfig("hailuo-test"))

	require.NoError(t, DeletePricingConfig("hailuo-test"))
	require.Nil(t, GetPricingConfig("hailuo-test"))
}