Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

243 righe
6.6 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/common"
  5. "github.com/glebarez/sqlite"
  6. "github.com/stretchr/testify/assert"
  7. "github.com/stretchr/testify/require"
  8. "gorm.io/gorm"
  9. )
  10. func setupChannelPricingDB(t *testing.T) *gorm.DB {
  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. common.UsingSQLite = true
  19. common.RedisEnabled = false
  20. require.NoError(t, db.AutoMigrate(&ChannelPricing{}))
  21. t.Cleanup(func() {
  22. DB = origDB
  23. sqlDB.Close()
  24. })
  25. return db
  26. }
  27. func TestCacheWriteThrough_Insert(t *testing.T) {
  28. setupChannelPricingDB(t)
  29. cp := &ChannelPricing{
  30. ModelName: "test-insert-model",
  31. ChannelId: 9001,
  32. QuotaType: QuotaTypeByTokens,
  33. ModelRatio: 1.0,
  34. CacheRatio: 0.8,
  35. ImageRatio: 1.2,
  36. }
  37. require.NoError(t, cp.Insert())
  38. defer cp.Delete()
  39. found, ok := GetEffectivePricing("test-insert-model", 9001)
  40. assert.True(t, ok)
  41. assert.Equal(t, 0.8, found.CacheRatio)
  42. assert.Equal(t, 1.2, found.ImageRatio)
  43. }
  44. func TestCacheWriteThrough_Update(t *testing.T) {
  45. setupChannelPricingDB(t)
  46. cp := &ChannelPricing{
  47. ModelName: "test-update-model",
  48. ChannelId: 9002,
  49. QuotaType: QuotaTypeByTokens,
  50. ModelRatio: 1.0,
  51. }
  52. require.NoError(t, cp.Insert())
  53. defer cp.Delete()
  54. cp.CacheRatio = 0.9
  55. cp.AudioRatio = 1.5
  56. require.NoError(t, cp.Update())
  57. found, ok := GetEffectivePricing("test-update-model", 9002)
  58. assert.True(t, ok)
  59. assert.Equal(t, 0.9, found.CacheRatio)
  60. assert.Equal(t, 1.5, found.AudioRatio)
  61. }
  62. func TestCacheWriteThrough_Delete(t *testing.T) {
  63. setupChannelPricingDB(t)
  64. cp := &ChannelPricing{
  65. ModelName: "test-delete-model",
  66. ChannelId: 9003,
  67. QuotaType: QuotaTypeByTokens,
  68. ModelRatio: 1.0,
  69. }
  70. require.NoError(t, cp.Insert())
  71. require.NoError(t, cp.Delete())
  72. found, ok := GetEffectivePricing("test-delete-model", 9003)
  73. assert.False(t, ok)
  74. assert.Nil(t, found)
  75. }
  76. func TestExtendedFields_DefaultZero(t *testing.T) {
  77. setupChannelPricingDB(t)
  78. cp := &ChannelPricing{
  79. ModelName: "test-default-model",
  80. ChannelId: 9004,
  81. QuotaType: QuotaTypeByTokens,
  82. ModelRatio: 1.0,
  83. }
  84. require.NoError(t, cp.Insert())
  85. defer cp.Delete()
  86. found, ok := GetEffectivePricing("test-default-model", 9004)
  87. assert.True(t, ok)
  88. assert.Equal(t, 0.0, found.CacheRatio)
  89. assert.Equal(t, 0.0, found.CacheCreationRatio)
  90. assert.Equal(t, 0.0, found.ImageRatio)
  91. assert.Equal(t, 0.0, found.AudioRatio)
  92. assert.Equal(t, 0.0, found.AudioCompletionRatio)
  93. }
  94. func TestGetEffectivePricing_NotFound(t *testing.T) {
  95. setupChannelPricingDB(t)
  96. found, ok := GetEffectivePricing("nonexistent-model-xyz", 99999)
  97. assert.False(t, ok)
  98. assert.Nil(t, found)
  99. }
  100. // === 默认通道测试 ===
  101. func TestDefaultChannel_SetAndGet(t *testing.T) {
  102. setupChannelPricingDB(t)
  103. // 创建两条定价记录
  104. cp1 := &ChannelPricing{ModelName: "default-test-model", ChannelId: 100, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0}
  105. cp2 := &ChannelPricing{ModelName: "default-test-model", ChannelId: 200, QuotaType: QuotaTypeByTokens, ModelRatio: 2.0}
  106. require.NoError(t, cp1.Insert())
  107. require.NoError(t, cp2.Insert())
  108. t.Cleanup(func() { cp1.Delete(); cp2.Delete() })
  109. // 初始没有默认
  110. _, ok := GetDefaultChannelId("default-test-model")
  111. assert.False(t, ok)
  112. // 设置通道 100 为默认
  113. require.NoError(t, SetDefaultChannel("default-test-model", 100))
  114. id, ok := GetDefaultChannelId("default-test-model")
  115. assert.True(t, ok)
  116. assert.Equal(t, 100, id)
  117. // 切换默认到通道 200
  118. require.NoError(t, SetDefaultChannel("default-test-model", 200))
  119. id, ok = GetDefaultChannelId("default-test-model")
  120. assert.True(t, ok)
  121. assert.Equal(t, 200, id)
  122. // 验证旧的默认标记被清除
  123. found, _ := GetEffectivePricing("default-test-model", 100)
  124. assert.False(t, found.IsDefault)
  125. found, _ = GetEffectivePricing("default-test-model", 200)
  126. assert.True(t, found.IsDefault)
  127. }
  128. func TestDefaultChannel_Clear(t *testing.T) {
  129. setupChannelPricingDB(t)
  130. cp := &ChannelPricing{ModelName: "clear-test-model", ChannelId: 300, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0}
  131. require.NoError(t, cp.Insert())
  132. t.Cleanup(func() { cp.Delete() })
  133. require.NoError(t, SetDefaultChannel("clear-test-model", 300))
  134. _, ok := GetDefaultChannelId("clear-test-model")
  135. assert.True(t, ok)
  136. require.NoError(t, ClearDefaultChannel("clear-test-model"))
  137. _, ok = GetDefaultChannelId("clear-test-model")
  138. assert.False(t, ok)
  139. }
  140. func TestDefaultChannel_DeleteClearsCache(t *testing.T) {
  141. setupChannelPricingDB(t)
  142. cp := &ChannelPricing{ModelName: "delete-test-model", ChannelId: 400, QuotaType: QuotaTypeByTokens, ModelRatio: 1.0, IsDefault: true}
  143. require.NoError(t, cp.Insert())
  144. id, ok := GetDefaultChannelId("delete-test-model")
  145. assert.True(t, ok)
  146. assert.Equal(t, 400, id)
  147. // 删除记录后缓存应清除
  148. require.NoError(t, cp.Delete())
  149. _, ok = GetDefaultChannelId("delete-test-model")
  150. assert.False(t, ok)
  151. }
  152. func TestDefaultChannel_LoadCache(t *testing.T) {
  153. db := setupChannelPricingDB(t)
  154. // 直接插入数据(绕过缓存)
  155. db.Create(&ChannelPricing{ModelName: "load-model", ChannelId: 500, ModelRatio: 1.0, IsDefault: true})
  156. db.Create(&ChannelPricing{ModelName: "load-model", ChannelId: 501, ModelRatio: 2.0, IsDefault: false})
  157. db.Create(&ChannelPricing{ModelName: "other-model", ChannelId: 502, ModelRatio: 1.0, IsDefault: true})
  158. // 全量加载缓存
  159. LoadChannelPricingCache()
  160. id, ok := GetDefaultChannelId("load-model")
  161. assert.True(t, ok)
  162. assert.Equal(t, 500, id)
  163. id, ok = GetDefaultChannelId("other-model")
  164. assert.True(t, ok)
  165. assert.Equal(t, 502, id)
  166. // 无默认的模型
  167. _, ok = GetDefaultChannelId("nonexistent")
  168. assert.False(t, ok)
  169. }
  170. func TestDefaultChannel_InsertWithDefault(t *testing.T) {
  171. setupChannelPricingDB(t)
  172. cp := &ChannelPricing{ModelName: "insert-default-model", ChannelId: 600, ModelRatio: 1.0, IsDefault: true}
  173. require.NoError(t, cp.Insert())
  174. t.Cleanup(func() { cp.Delete() })
  175. id, ok := GetDefaultChannelId("insert-default-model")
  176. assert.True(t, ok)
  177. assert.Equal(t, 600, id)
  178. }
  179. func TestDefaultChannel_UpdateWithDefault(t *testing.T) {
  180. setupChannelPricingDB(t)
  181. cp1 := &ChannelPricing{ModelName: "update-default-model", ChannelId: 700, ModelRatio: 1.0, IsDefault: true}
  182. cp2 := &ChannelPricing{ModelName: "update-default-model", ChannelId: 701, ModelRatio: 2.0}
  183. require.NoError(t, cp1.Insert())
  184. require.NoError(t, cp2.Insert())
  185. t.Cleanup(func() { cp1.Delete(); cp2.Delete() })
  186. // cp1 是默认,通过 Update 把 cp2 设为默认
  187. cp2.IsDefault = true
  188. require.NoError(t, cp2.Update())
  189. id, ok := GetDefaultChannelId("update-default-model")
  190. assert.True(t, ok)
  191. assert.Equal(t, 701, id)
  192. }