|
- package helper
-
- import (
- "net/http"
- "net/http/httptest"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/dto"
- "github.com/QuantumNous/new-api/model"
- relaycommon "github.com/QuantumNous/new-api/relay/common"
- "github.com/QuantumNous/new-api/setting/ratio_setting"
- "github.com/QuantumNous/new-api/types"
- "github.com/glebarez/sqlite"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- const testModelUsePriceSwitch = "test-useprice-switch-model"
-
- func setupUsePriceSwitchTest(t *testing.T) {
- t.Helper()
-
- // 初始化 SQLite 内存数据库
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := model.DB
- model.DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&model.ChannelPricing{}))
-
- // 全局定价:按次计费(modelPrice 存在 → usePrice=true)
- require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testModelUsePriceSwitch+`":0.5}`))
- // 全局 modelRatio 也设上(渠道按量计费时会用到)
- require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testModelUsePriceSwitch+`":15}`))
- require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3}`))
- // 全局扩展比率
- require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":0.1}`))
- require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":1.25}`))
- require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{"`+testModelUsePriceSwitch+`":2.0}`))
- require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{"`+testModelUsePriceSwitch+`":5.0}`))
- require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3.0}`))
-
- // 渠道定价:按量计费(QuotaTypeByTokens),扩展比率全部为 0(用全局)
- cp := &model.ChannelPricing{
- ModelName: testModelUsePriceSwitch,
- ChannelId: 9901,
- QuotaType: model.QuotaTypeByTokens,
- ModelRatio: 15,
- CompletionRatio: 3,
- // 扩展比率全部为 0 → 意味着回退全局
- }
- require.NoError(t, cp.Insert())
-
- t.Cleanup(func() {
- model.DB = origDB
- sqlDB.Close()
- // 清理全局定价
- ratio_setting.UpdateModelPriceByJSONString(`{}`)
- ratio_setting.UpdateModelRatioByJSONString(`{}`)
- ratio_setting.UpdateCompletionRatioByJSONString(`{}`)
- ratio_setting.UpdateCacheRatioByJSONString(`{}`)
- ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)
- ratio_setting.UpdateImageRatioByJSONString(`{}`)
- ratio_setting.UpdateAudioRatioByJSONString(`{}`)
- ratio_setting.UpdateAudioCompletionRatioByJSONString(`{}`)
- })
- }
-
- func buildTestContext(t *testing.T) *gin.Context {
- t.Helper()
- gin.SetMode(gin.TestMode)
- w := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(w)
- c.Request, _ = http.NewRequest("POST", "/v1/chat/completions", nil)
- return c
- }
-
- // TestGlobalRatiosFallbackWhenGlobalUsePriceTrue 验证核心场景:
- // 全局按次计费 → 渠道按量计费 → PriceData 中扩展比率仍应被全局默认值填充
- // 而不是保持 0(否则后续 UpdatePriceDataForChannelPricing 切换为按量时会丢失比率)
- func TestGlobalRatiosFallbackWhenGlobalUsePriceTrue(t *testing.T) {
- setupUsePriceSwitchTest(t)
- c := buildTestContext(t)
-
- // 模拟渠道选择后的场景:ChannelMeta 不为空,能查到渠道定价
- info := &relaycommon.RelayInfo{
- OriginModelName: testModelUsePriceSwitch,
- ChannelMeta: &relaycommon.ChannelMeta{
- ChannelId: 9901,
- },
- UserSetting: dto.UserSetting{},
- }
- meta := &types.TokenCountMeta{}
-
- priceData, err := ModelPriceHelper(c, info, 100, meta)
- require.NoError(t, err)
-
- // 断言:渠道按量计费覆盖了全局按次,扩展比率必须被全局默认值填充,不能是 0
- assert.Equal(t, 0.1, priceData.CacheRatio, "CacheRatio 应从全局回退填充")
- assert.Equal(t, 1.25, priceData.CacheCreationRatio, "CacheCreationRatio 应从全局回退填充")
- assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 应从全局回退填充")
- assert.Equal(t, 5.0, priceData.AudioRatio, "AudioRatio 应从全局回退填充")
- assert.Equal(t, 3.0, priceData.AudioCompletionRatio, "AudioCompletionRatio 应从全局回退填充")
- }
-
- // TestGlobalRatiosFallbackWithoutChannelPricing 验证无渠道定价时全局比率也正常回退
- func TestGlobalRatiosFallbackWithoutChannelPricing(t *testing.T) {
- setupUsePriceSwitchTest(t)
- c := buildTestContext(t)
-
- // 无 ChannelMeta → 渠道定价不可用,走全局
- info := &relaycommon.RelayInfo{
- OriginModelName: testModelUsePriceSwitch,
- UserSetting: dto.UserSetting{},
- }
- meta := &types.TokenCountMeta{}
-
- priceData, err := ModelPriceHelper(c, info, 100, meta)
- require.NoError(t, err)
-
- // 全局是按次计费,UsePrice=true,但比率仍应被填充
- assert.True(t, priceData.UsePrice)
- assert.Equal(t, 0.1, priceData.CacheRatio, "CacheRatio 应从全局回退填充")
- assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 应从全局回退填充")
- }
-
- // TestChannelRatiosOverrideGlobal 验证渠道非零值优先于全局
- func TestChannelRatiosOverrideGlobal(t *testing.T) {
- setupUsePriceSwitchTest(t)
-
- // 渠道定价设了自定义 cacheRatio=0.5
- cp := &model.ChannelPricing{
- ModelName: testModelUsePriceSwitch,
- ChannelId: 9902,
- QuotaType: model.QuotaTypeByTokens,
- ModelRatio: 15,
- CompletionRatio: 3,
- CacheRatio: 0.5, // 渠道自定义,覆盖全局 0.1
- // 其他扩展比率保持 0,用全局
- }
- require.NoError(t, cp.Insert())
-
- c := buildTestContext(t)
- info := &relaycommon.RelayInfo{
- OriginModelName: testModelUsePriceSwitch,
- ChannelMeta: &relaycommon.ChannelMeta{
- ChannelId: 9902,
- },
- UserSetting: dto.UserSetting{},
- }
-
- priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
- require.NoError(t, err)
-
- assert.Equal(t, 0.5, priceData.CacheRatio, "渠道自定义 CacheRatio 应覆盖全局")
- assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 未设渠道值,应从全局回退")
- assert.Equal(t, 5.0, priceData.AudioRatio, "AudioRatio 未设渠道值,应从全局回退")
- }
|