You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

167 lines
6.2 KiB

  1. package helper
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "testing"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/dto"
  8. "github.com/QuantumNous/new-api/model"
  9. relaycommon "github.com/QuantumNous/new-api/relay/common"
  10. "github.com/QuantumNous/new-api/setting/ratio_setting"
  11. "github.com/QuantumNous/new-api/types"
  12. "github.com/glebarez/sqlite"
  13. "github.com/gin-gonic/gin"
  14. "github.com/stretchr/testify/assert"
  15. "github.com/stretchr/testify/require"
  16. "gorm.io/gorm"
  17. )
  18. const testModelUsePriceSwitch = "test-useprice-switch-model"
  19. func setupUsePriceSwitchTest(t *testing.T) {
  20. t.Helper()
  21. // 初始化 SQLite 内存数据库
  22. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  23. require.NoError(t, err)
  24. sqlDB, _ := db.DB()
  25. sqlDB.SetMaxOpenConns(1)
  26. origDB := model.DB
  27. model.DB = db
  28. common.UsingSQLite = true
  29. common.RedisEnabled = false
  30. require.NoError(t, db.AutoMigrate(&model.ChannelPricing{}))
  31. // 全局定价:按次计费(modelPrice 存在 → usePrice=true)
  32. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testModelUsePriceSwitch+`":0.5}`))
  33. // 全局 modelRatio 也设上(渠道按量计费时会用到)
  34. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testModelUsePriceSwitch+`":15}`))
  35. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3}`))
  36. // 全局扩展比率
  37. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":0.1}`))
  38. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":1.25}`))
  39. require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{"`+testModelUsePriceSwitch+`":2.0}`))
  40. require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{"`+testModelUsePriceSwitch+`":5.0}`))
  41. require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3.0}`))
  42. // 渠道定价:按量计费(QuotaTypeByTokens),扩展比率全部为 0(用全局)
  43. cp := &model.ChannelPricing{
  44. ModelName: testModelUsePriceSwitch,
  45. ChannelId: 9901,
  46. QuotaType: model.QuotaTypeByTokens,
  47. ModelRatio: 15,
  48. CompletionRatio: 3,
  49. // 扩展比率全部为 0 → 意味着回退全局
  50. }
  51. require.NoError(t, cp.Insert())
  52. t.Cleanup(func() {
  53. model.DB = origDB
  54. sqlDB.Close()
  55. // 清理全局定价
  56. ratio_setting.UpdateModelPriceByJSONString(`{}`)
  57. ratio_setting.UpdateModelRatioByJSONString(`{}`)
  58. ratio_setting.UpdateCompletionRatioByJSONString(`{}`)
  59. ratio_setting.UpdateCacheRatioByJSONString(`{}`)
  60. ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)
  61. ratio_setting.UpdateImageRatioByJSONString(`{}`)
  62. ratio_setting.UpdateAudioRatioByJSONString(`{}`)
  63. ratio_setting.UpdateAudioCompletionRatioByJSONString(`{}`)
  64. })
  65. }
  66. func buildTestContext(t *testing.T) *gin.Context {
  67. t.Helper()
  68. gin.SetMode(gin.TestMode)
  69. w := httptest.NewRecorder()
  70. c, _ := gin.CreateTestContext(w)
  71. c.Request, _ = http.NewRequest("POST", "/v1/chat/completions", nil)
  72. return c
  73. }
  74. // TestGlobalRatiosFallbackWhenGlobalUsePriceTrue 验证核心场景:
  75. // 全局按次计费 → 渠道按量计费 → PriceData 中扩展比率仍应被全局默认值填充
  76. // 而不是保持 0(否则后续 UpdatePriceDataForChannelPricing 切换为按量时会丢失比率)
  77. func TestGlobalRatiosFallbackWhenGlobalUsePriceTrue(t *testing.T) {
  78. setupUsePriceSwitchTest(t)
  79. c := buildTestContext(t)
  80. // 模拟渠道选择后的场景:ChannelMeta 不为空,能查到渠道定价
  81. info := &relaycommon.RelayInfo{
  82. OriginModelName: testModelUsePriceSwitch,
  83. ChannelMeta: &relaycommon.ChannelMeta{
  84. ChannelId: 9901,
  85. },
  86. UserSetting: dto.UserSetting{},
  87. }
  88. meta := &types.TokenCountMeta{}
  89. priceData, err := ModelPriceHelper(c, info, 100, meta)
  90. require.NoError(t, err)
  91. // 断言:渠道按量计费覆盖了全局按次,扩展比率必须被全局默认值填充,不能是 0
  92. assert.Equal(t, 0.1, priceData.CacheRatio, "CacheRatio 应从全局回退填充")
  93. assert.Equal(t, 1.25, priceData.CacheCreationRatio, "CacheCreationRatio 应从全局回退填充")
  94. assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 应从全局回退填充")
  95. assert.Equal(t, 5.0, priceData.AudioRatio, "AudioRatio 应从全局回退填充")
  96. assert.Equal(t, 3.0, priceData.AudioCompletionRatio, "AudioCompletionRatio 应从全局回退填充")
  97. }
  98. // TestGlobalRatiosFallbackWithoutChannelPricing 验证无渠道定价时全局比率也正常回退
  99. func TestGlobalRatiosFallbackWithoutChannelPricing(t *testing.T) {
  100. setupUsePriceSwitchTest(t)
  101. c := buildTestContext(t)
  102. // 无 ChannelMeta → 渠道定价不可用,走全局
  103. info := &relaycommon.RelayInfo{
  104. OriginModelName: testModelUsePriceSwitch,
  105. UserSetting: dto.UserSetting{},
  106. }
  107. meta := &types.TokenCountMeta{}
  108. priceData, err := ModelPriceHelper(c, info, 100, meta)
  109. require.NoError(t, err)
  110. // 全局是按次计费,UsePrice=true,但比率仍应被填充
  111. assert.True(t, priceData.UsePrice)
  112. assert.Equal(t, 0.1, priceData.CacheRatio, "CacheRatio 应从全局回退填充")
  113. assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 应从全局回退填充")
  114. }
  115. // TestChannelRatiosOverrideGlobal 验证渠道非零值优先于全局
  116. func TestChannelRatiosOverrideGlobal(t *testing.T) {
  117. setupUsePriceSwitchTest(t)
  118. // 渠道定价设了自定义 cacheRatio=0.5
  119. cp := &model.ChannelPricing{
  120. ModelName: testModelUsePriceSwitch,
  121. ChannelId: 9902,
  122. QuotaType: model.QuotaTypeByTokens,
  123. ModelRatio: 15,
  124. CompletionRatio: 3,
  125. CacheRatio: 0.5, // 渠道自定义,覆盖全局 0.1
  126. // 其他扩展比率保持 0,用全局
  127. }
  128. require.NoError(t, cp.Insert())
  129. c := buildTestContext(t)
  130. info := &relaycommon.RelayInfo{
  131. OriginModelName: testModelUsePriceSwitch,
  132. ChannelMeta: &relaycommon.ChannelMeta{
  133. ChannelId: 9902,
  134. },
  135. UserSetting: dto.UserSetting{},
  136. }
  137. priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
  138. require.NoError(t, err)
  139. assert.Equal(t, 0.5, priceData.CacheRatio, "渠道自定义 CacheRatio 应覆盖全局")
  140. assert.Equal(t, 2.0, priceData.ImageRatio, "ImageRatio 未设渠道值,应从全局回退")
  141. assert.Equal(t, 5.0, priceData.AudioRatio, "AudioRatio 未设渠道值,应从全局回退")
  142. }