|
- 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, "应选最便宜的按次渠道")
- }
|