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.
 
 
 

123 line
4.0 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/setting/ratio_setting"
  5. "github.com/glebarez/sqlite"
  6. "github.com/stretchr/testify/require"
  7. "gorm.io/gorm"
  8. )
  9. const testPricingModel = "test-pricing-model-apply"
  10. func setupPricingTest(t *testing.T) {
  11. t.Helper()
  12. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  13. require.NoError(t, err)
  14. sqlDB, _ := db.DB()
  15. sqlDB.SetMaxOpenConns(1)
  16. origDB := DB
  17. DB = db
  18. require.NoError(t, db.AutoMigrate(&ChannelPricing{}))
  19. // 全局默认定价
  20. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`))
  21. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`))
  22. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`))
  23. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`))
  24. t.Cleanup(func() {
  25. DB = origDB
  26. sqlDB.Close()
  27. ratio_setting.UpdateModelRatioByJSONString(`{}`)
  28. ratio_setting.UpdateCompletionRatioByJSONString(`{}`)
  29. ratio_setting.UpdateCacheRatioByJSONString(`{}`)
  30. ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)
  31. ratio_setting.UpdateModelPriceByJSONString(`{}`)
  32. })
  33. }
  34. func TestApplyBestChannelPricing_NoChannelPricing(t *testing.T) {
  35. setupPricingTest(t)
  36. p := &Pricing{}
  37. applyBestChannelPricing(p, nil, testPricingModel)
  38. require.Equal(t, 0, p.QuotaType, "应使用全局按量计费")
  39. require.Equal(t, 10.0, p.ModelRatio, "应使用全局 model_ratio")
  40. require.Equal(t, 3.0, p.CompletionRatio, "应使用全局 completion_ratio")
  41. require.Equal(t, 0.5, p.CacheRatio, "应使用全局 cache_ratio")
  42. require.Equal(t, 0.75, p.CacheCreationRatio, "应使用全局 cache_creation_ratio")
  43. }
  44. func TestApplyBestChannelPricing_PerTokenCheapest(t *testing.T) {
  45. setupPricingTest(t)
  46. // 渠道A: model_ratio=8(更便宜)
  47. require.NoError(t, (&ChannelPricing{
  48. ModelName: testPricingModel, ChannelId: 1,
  49. QuotaType: QuotaTypeByTokens, ModelRatio: 8, CompletionRatio: 2,
  50. CacheRatio: 0.3, CacheCreationRatio: 0.6,
  51. }).Insert())
  52. // 渠道B: model_ratio=12(更贵)
  53. require.NoError(t, (&ChannelPricing{
  54. ModelName: testPricingModel, ChannelId: 2,
  55. QuotaType: QuotaTypeByTokens, ModelRatio: 12, CompletionRatio: 4,
  56. CacheRatio: 0.8, CacheCreationRatio: 1.0,
  57. }).Insert())
  58. cps := []ChannelPricing{
  59. {ModelRatio: 8, CompletionRatio: 2, CacheRatio: 0.3, CacheCreationRatio: 0.6},
  60. {ModelRatio: 12, CompletionRatio: 4, CacheRatio: 0.8, CacheCreationRatio: 1.0},
  61. }
  62. p := &Pricing{}
  63. applyBestChannelPricing(p, cps, testPricingModel)
  64. require.Equal(t, 0, p.QuotaType)
  65. require.Equal(t, 8.0, p.ModelRatio, "应选最便宜的渠道A")
  66. require.Equal(t, 2.0, p.CompletionRatio)
  67. require.Equal(t, 0.3, p.CacheRatio)
  68. require.Equal(t, 0.6, p.CacheCreationRatio)
  69. }
  70. func TestApplyBestChannelPricing_ExtendedRatioZeroFallback(t *testing.T) {
  71. setupPricingTest(t)
  72. // 渠道定价中扩展比率为 0,应回退到全局值
  73. require.NoError(t, (&ChannelPricing{
  74. ModelName: testPricingModel, ChannelId: 1,
  75. QuotaType: QuotaTypeByTokens, ModelRatio: 5, CompletionRatio: 1,
  76. CacheRatio: 0, CacheCreationRatio: 0, // 0 = 未设置
  77. }).Insert())
  78. cps := []ChannelPricing{
  79. {ModelRatio: 5, CompletionRatio: 1, CacheRatio: 0, CacheCreationRatio: 0},
  80. }
  81. p := &Pricing{}
  82. applyBestChannelPricing(p, cps, testPricingModel)
  83. require.Equal(t, 0.5, p.CacheRatio, "cache_ratio=0 应回退全局 0.5")
  84. require.Equal(t, 0.75, p.CacheCreationRatio, "cache_creation_ratio=0 应回退全局 0.75")
  85. }
  86. func TestApplyBestChannelPricing_PerCallCheapest(t *testing.T) {
  87. setupPricingTest(t)
  88. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingModel+`":0.8}`))
  89. // 只有按次计费的渠道
  90. cps := []ChannelPricing{
  91. {QuotaType: QuotaTypeByCall, ModelPrice: 0.3},
  92. {QuotaType: QuotaTypeByCall, ModelPrice: 0.5},
  93. }
  94. p := &Pricing{}
  95. applyBestChannelPricing(p, cps, testPricingModel)
  96. require.Equal(t, 1, p.QuotaType)
  97. require.Equal(t, 0.3, p.ModelPrice, "应选最便宜的按次渠道")
  98. }