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")) }