From 415174a676881264fd86322e5ef55669b52a1295 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Fri, 29 May 2026 14:53:09 +0800 Subject: [PATCH] test(model): switch pricing tests to global defaults --- model/channel_pricing_test.go | 121 ---------------------------------- model/pricing_test.go | 119 +++++++-------------------------- 2 files changed, 23 insertions(+), 217 deletions(-) delete mode 100644 model/channel_pricing_test.go diff --git a/model/channel_pricing_test.go b/model/channel_pricing_test.go deleted file mode 100644 index 9ed8a9f..0000000 --- a/model/channel_pricing_test.go +++ /dev/null @@ -1,121 +0,0 @@ -package model - -import ( - "testing" - - "github.com/QuantumNous/new-api/common" - "github.com/glebarez/sqlite" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "gorm.io/gorm" -) - -func setupChannelPricingDB(t *testing.T) *gorm.DB { - 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 - common.UsingSQLite = true - common.RedisEnabled = false - - require.NoError(t, db.AutoMigrate(&ChannelPricing{})) - - t.Cleanup(func() { - DB = origDB - sqlDB.Close() - }) - return db -} - -func TestCacheWriteThrough_Insert(t *testing.T) { - setupChannelPricingDB(t) - - cp := &ChannelPricing{ - ModelName: "test-insert-model", - ChannelId: 9001, - QuotaType: QuotaTypeByTokens, - ModelRatio: 1.0, - CacheRatio: 0.8, - ImageRatio: 1.2, - } - require.NoError(t, cp.Insert()) - defer cp.Delete() - - found, ok := GetEffectivePricing("test-insert-model", 9001) - assert.True(t, ok) - assert.Equal(t, 0.8, found.CacheRatio) - assert.Equal(t, 1.2, found.ImageRatio) -} - -func TestCacheWriteThrough_Update(t *testing.T) { - setupChannelPricingDB(t) - - cp := &ChannelPricing{ - ModelName: "test-update-model", - ChannelId: 9002, - QuotaType: QuotaTypeByTokens, - ModelRatio: 1.0, - } - require.NoError(t, cp.Insert()) - defer cp.Delete() - - cp.CacheRatio = 0.9 - cp.AudioRatio = 1.5 - require.NoError(t, cp.Update()) - - found, ok := GetEffectivePricing("test-update-model", 9002) - assert.True(t, ok) - assert.Equal(t, 0.9, found.CacheRatio) - assert.Equal(t, 1.5, found.AudioRatio) -} - -func TestCacheWriteThrough_Delete(t *testing.T) { - setupChannelPricingDB(t) - - cp := &ChannelPricing{ - ModelName: "test-delete-model", - ChannelId: 9003, - QuotaType: QuotaTypeByTokens, - ModelRatio: 1.0, - } - require.NoError(t, cp.Insert()) - - require.NoError(t, cp.Delete()) - - found, ok := GetEffectivePricing("test-delete-model", 9003) - assert.False(t, ok) - assert.Nil(t, found) -} - -func TestExtendedFields_DefaultZero(t *testing.T) { - setupChannelPricingDB(t) - - cp := &ChannelPricing{ - ModelName: "test-default-model", - ChannelId: 9004, - QuotaType: QuotaTypeByTokens, - ModelRatio: 1.0, - } - require.NoError(t, cp.Insert()) - defer cp.Delete() - - found, ok := GetEffectivePricing("test-default-model", 9004) - assert.True(t, ok) - assert.Equal(t, 0.0, found.CacheRatio) - assert.Equal(t, 0.0, found.CacheCreationRatio) - assert.Equal(t, 0.0, found.ImageRatio) - assert.Equal(t, 0.0, found.AudioRatio) - assert.Equal(t, 0.0, found.AudioCompletionRatio) -} - -func TestGetEffectivePricing_NotFound(t *testing.T) { - setupChannelPricingDB(t) - - found, ok := GetEffectivePricing("nonexistent-model-xyz", 99999) - assert.False(t, ok) - assert.Nil(t, found) -} diff --git a/model/pricing_test.go b/model/pricing_test.go index 6364d53..749b66b 100644 --- a/model/pricing_test.go +++ b/model/pricing_test.go @@ -4,119 +4,46 @@ 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" +const testPricingCallModel = "test-pricing-call-model" -func setupPricingTest(t *testing.T) { +func resetPricingRatios(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(`{}`)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`)) + require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`)) + require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) +} - // 全局默认定价 +func TestApplyGlobalDefault_PerTokenModel(t *testing.T) { + resetPricingRatios(t) 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) + applyGlobalDefault(p, 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") + require.Equal(t, QuotaTypeByTokens, p.QuotaType) + require.Equal(t, 10.0, p.ModelRatio) + require.Equal(t, 3.0, p.CompletionRatio) + require.Equal(t, 0.5, p.CacheRatio) + require.Equal(t, 0.75, p.CacheCreationRatio) } -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}, - } +func TestApplyGlobalDefault_PerCallModel(t *testing.T) { + resetPricingRatios(t) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`)) p := &Pricing{} - applyBestChannelPricing(p, cps, testPricingModel) + applyGlobalDefault(p, testPricingCallModel) - require.Equal(t, 1, p.QuotaType) - require.Equal(t, 0.3, p.ModelPrice, "应选最便宜的按次渠道") + require.Equal(t, QuotaTypeByCall, p.QuotaType) + require.Equal(t, 0.3, p.ModelPrice) + require.Equal(t, 0.0, p.ModelRatio) }