package model import ( "testing" "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" "gorm.io/gorm" ) const testPricingModel = "test-pricing-model-apply" func setupPricingTest(t *testing.T) { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) sqlDB, _ := db.DB() sqlDB.SetMaxOpenConns(1) origDB := DB DB = db require.NoError(t, db.AutoMigrate(&ChannelPricing{})) // 全局默认定价 require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`)) require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`)) require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`)) require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`)) t.Cleanup(func() { DB = origDB sqlDB.Close() ratio_setting.UpdateModelRatioByJSONString(`{}`) ratio_setting.UpdateCompletionRatioByJSONString(`{}`) ratio_setting.UpdateCacheRatioByJSONString(`{}`) ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`) ratio_setting.UpdateModelPriceByJSONString(`{}`) }) } func TestApplyBestChannelPricing_NoChannelPricing(t *testing.T) { setupPricingTest(t) p := &Pricing{} applyBestChannelPricing(p, nil, testPricingModel) require.Equal(t, 0, p.QuotaType, "应使用全局按量计费") require.Equal(t, 10.0, p.ModelRatio, "应使用全局 model_ratio") require.Equal(t, 3.0, p.CompletionRatio, "应使用全局 completion_ratio") require.Equal(t, 0.5, p.CacheRatio, "应使用全局 cache_ratio") require.Equal(t, 0.75, p.CacheCreationRatio, "应使用全局 cache_creation_ratio") } func TestApplyBestChannelPricing_PerTokenCheapest(t *testing.T) { setupPricingTest(t) // 渠道A: model_ratio=8(更便宜) require.NoError(t, (&ChannelPricing{ ModelName: testPricingModel, ChannelId: 1, QuotaType: QuotaTypeByTokens, ModelRatio: 8, CompletionRatio: 2, CacheRatio: 0.3, CacheCreationRatio: 0.6, }).Insert()) // 渠道B: model_ratio=12(更贵) require.NoError(t, (&ChannelPricing{ ModelName: testPricingModel, ChannelId: 2, QuotaType: QuotaTypeByTokens, ModelRatio: 12, CompletionRatio: 4, CacheRatio: 0.8, CacheCreationRatio: 1.0, }).Insert()) cps := []ChannelPricing{ {ModelRatio: 8, CompletionRatio: 2, CacheRatio: 0.3, CacheCreationRatio: 0.6}, {ModelRatio: 12, CompletionRatio: 4, CacheRatio: 0.8, CacheCreationRatio: 1.0}, } p := &Pricing{} applyBestChannelPricing(p, cps, testPricingModel) require.Equal(t, 0, p.QuotaType) require.Equal(t, 8.0, p.ModelRatio, "应选最便宜的渠道A") require.Equal(t, 2.0, p.CompletionRatio) require.Equal(t, 0.3, p.CacheRatio) require.Equal(t, 0.6, p.CacheCreationRatio) } func TestApplyBestChannelPricing_ExtendedRatioZeroFallback(t *testing.T) { setupPricingTest(t) // 渠道定价中扩展比率为 0,应回退到全局值 require.NoError(t, (&ChannelPricing{ ModelName: testPricingModel, ChannelId: 1, QuotaType: QuotaTypeByTokens, ModelRatio: 5, CompletionRatio: 1, CacheRatio: 0, CacheCreationRatio: 0, // 0 = 未设置 }).Insert()) cps := []ChannelPricing{ {ModelRatio: 5, CompletionRatio: 1, CacheRatio: 0, CacheCreationRatio: 0}, } p := &Pricing{} applyBestChannelPricing(p, cps, testPricingModel) require.Equal(t, 0.5, p.CacheRatio, "cache_ratio=0 应回退全局 0.5") require.Equal(t, 0.75, p.CacheCreationRatio, "cache_creation_ratio=0 应回退全局 0.75") } func TestApplyBestChannelPricing_PerCallCheapest(t *testing.T) { setupPricingTest(t) require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingModel+`":0.8}`)) // 只有按次计费的渠道 cps := []ChannelPricing{ {QuotaType: QuotaTypeByCall, ModelPrice: 0.3}, {QuotaType: QuotaTypeByCall, ModelPrice: 0.5}, } p := &Pricing{} applyBestChannelPricing(p, cps, testPricingModel) require.Equal(t, 1, p.QuotaType) require.Equal(t, 0.3, p.ModelPrice, "应选最便宜的按次渠道") }