Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.
 
 
 

284 рядки
7.8 KiB

  1. package controller
  2. import (
  3. "encoding/json"
  4. "net/http"
  5. "net/http/httptest"
  6. "testing"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/model"
  9. "github.com/QuantumNous/new-api/setting/ratio_setting"
  10. "github.com/glebarez/sqlite"
  11. "github.com/gin-gonic/gin"
  12. "github.com/stretchr/testify/assert"
  13. "github.com/stretchr/testify/require"
  14. "gorm.io/gorm"
  15. )
  16. // setupPricingTestDB 初始化测试数据库
  17. func setupPricingTestDB(t *testing.T) *gorm.DB {
  18. t.Helper()
  19. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  20. require.NoError(t, err)
  21. sqlDB, _ := db.DB()
  22. sqlDB.SetMaxOpenConns(1)
  23. origDB := model.DB
  24. model.DB = db
  25. common.UsingSQLite = true
  26. common.RedisEnabled = false
  27. require.NoError(t, db.AutoMigrate(&model.User{}))
  28. t.Cleanup(func() {
  29. model.DB = origDB
  30. sqlDB.Close()
  31. })
  32. return db
  33. }
  34. // setupPricingTestRouter 创建无认证路由(未登录场景)
  35. func setupPricingTestRouter() *gin.Engine {
  36. gin.SetMode(gin.TestMode)
  37. r := gin.New()
  38. r.GET("/api/pricing/user/:model", GetUserPricing)
  39. return r
  40. }
  41. // setupAuthRouter 创建带用户 ID 注入的路由(已登录场景)
  42. func setupAuthRouter(userID int) *gin.Engine {
  43. gin.SetMode(gin.TestMode)
  44. r := gin.New()
  45. r.GET("/api/pricing/user/:model", func(c *gin.Context) {
  46. c.Set("id", userID)
  47. c.Next()
  48. }, GetUserPricing)
  49. return r
  50. }
  51. // setPricingCache 直接设置定价缓存用于测试
  52. func setPricingCache(pricing []model.Pricing) {
  53. model.SetTestPricing(pricing)
  54. }
  55. // withGroupRatio 临时设置分组倍率,测试结束后恢复
  56. func withGroupRatio(t *testing.T, jsonStr string) {
  57. t.Helper()
  58. original := ratio_setting.GetGroupRatioCopy()
  59. ratio_setting.UpdateGroupRatioByJSONString(jsonStr)
  60. t.Cleanup(func() {
  61. origJSON, _ := json.Marshal(original)
  62. ratio_setting.UpdateGroupRatioByJSONString(string(origJSON))
  63. })
  64. }
  65. // ---- 测试用例 ----
  66. // TestGetUserPricing_ModelNotFound 模型不存在时应返回错误
  67. func TestGetUserPricing_ModelNotFound(t *testing.T) {
  68. setupPricingTestDB(t)
  69. router := setupPricingTestRouter()
  70. setPricingCache([]model.Pricing{})
  71. w := httptest.NewRecorder()
  72. req, _ := http.NewRequest("GET", "/api/pricing/user/nonexistent-model", nil)
  73. router.ServeHTTP(w, req)
  74. var resp map[string]interface{}
  75. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  76. assert.False(t, resp["success"].(bool))
  77. assert.Contains(t, resp["message"], "未找到")
  78. }
  79. // TestGetUserPricing_NotLoggedIn 未登录用户应只返回原价
  80. func TestGetUserPricing_NotLoggedIn(t *testing.T) {
  81. setupPricingTestDB(t)
  82. router := setupPricingTestRouter()
  83. setPricingCache([]model.Pricing{
  84. {
  85. ModelName: "gpt-4o",
  86. QuotaType: 0,
  87. ModelRatio: 15,
  88. CompletionRatio: 4,
  89. EnableGroup: []string{"default", "vip"},
  90. },
  91. })
  92. w := httptest.NewRecorder()
  93. req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
  94. router.ServeHTTP(w, req)
  95. var resp map[string]interface{}
  96. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  97. assert.True(t, resp["success"].(bool))
  98. assert.Equal(t, false, resp["logged_in"])
  99. assert.Equal(t, "gpt-4o", resp["model_name"])
  100. assert.Equal(t, float64(0), resp["quota_type"])
  101. // 验证原价: model_ratio * 2 = 15 * 2 = 30
  102. assert.Equal(t, float64(30), resp["original_input"])
  103. // 输出原价: model_ratio * completion_ratio * 2 = 15 * 4 * 2 = 120
  104. assert.Equal(t, float64(120), resp["original_output"])
  105. // 不应有用户价字段
  106. _, hasUserInput := resp["user_input"]
  107. assert.False(t, hasUserInput)
  108. }
  109. // TestGetUserPricing_NotLoggedIn_PerCall 按次计费模型,未登录
  110. func TestGetUserPricing_NotLoggedIn_PerCall(t *testing.T) {
  111. setupPricingTestDB(t)
  112. router := setupPricingTestRouter()
  113. setPricingCache([]model.Pricing{
  114. {
  115. ModelName: "dall-e-3",
  116. QuotaType: 1,
  117. ModelPrice: 0.04,
  118. EnableGroup: []string{"default"},
  119. },
  120. })
  121. w := httptest.NewRecorder()
  122. req, _ := http.NewRequest("GET", "/api/pricing/user/dall-e-3", nil)
  123. router.ServeHTTP(w, req)
  124. var resp map[string]interface{}
  125. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  126. assert.True(t, resp["success"].(bool))
  127. assert.Equal(t, false, resp["logged_in"])
  128. assert.Equal(t, float64(0.04), resp["original_price"])
  129. }
  130. // TestGetUserPricing_LoggedIn_NoDiscount 已登录但无折扣(分组倍率=1,无个人倍率)
  131. func TestGetUserPricing_LoggedIn_NoDiscount(t *testing.T) {
  132. db := setupPricingTestDB(t)
  133. setPricingCache([]model.Pricing{
  134. {
  135. ModelName: "gpt-4o",
  136. QuotaType: 0,
  137. ModelRatio: 15,
  138. CompletionRatio: 4,
  139. EnableGroup: []string{"default", "vip"},
  140. },
  141. })
  142. // 创建测试用户(default 分组,默认倍率为 1)
  143. user := &model.User{Id: 100, Group: "default", Username: "testuser", Status: 1}
  144. require.NoError(t, db.Create(user).Error)
  145. router := setupAuthRouter(100)
  146. w := httptest.NewRecorder()
  147. req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
  148. router.ServeHTTP(w, req)
  149. var resp map[string]interface{}
  150. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  151. assert.True(t, resp["success"].(bool))
  152. assert.Equal(t, true, resp["logged_in"])
  153. // default 分组默认倍率为 1,无折扣
  154. assert.Equal(t, float64(0), resp["savings_percent"])
  155. }
  156. // TestGetUserPricing_LoggedIn_GroupDiscount 已登录,有分组折扣
  157. func TestGetUserPricing_LoggedIn_GroupDiscount(t *testing.T) {
  158. db := setupPricingTestDB(t)
  159. setPricingCache([]model.Pricing{
  160. {
  161. ModelName: "gpt-4o",
  162. QuotaType: 0,
  163. ModelRatio: 15,
  164. CompletionRatio: 4,
  165. EnableGroup: []string{"default", "vip"},
  166. },
  167. })
  168. // 创建 VIP 用户
  169. user := &model.User{Id: 200, Group: "vip", Username: "vipuser", Status: 1}
  170. require.NoError(t, db.Create(user).Error)
  171. // 设置 VIP 分组倍率为 0.8
  172. withGroupRatio(t, `{"default":1,"vip":0.8}`)
  173. router := setupAuthRouter(200)
  174. w := httptest.NewRecorder()
  175. req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
  176. router.ServeHTTP(w, req)
  177. var resp map[string]interface{}
  178. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  179. assert.True(t, resp["success"].(bool))
  180. assert.Equal(t, true, resp["logged_in"])
  181. assert.Equal(t, "vip", resp["group"])
  182. assert.Equal(t, float64(0.8), resp["group_ratio"])
  183. // 用户价 = 原价 * 0.8
  184. // 输入: 30 * 0.8 = 24
  185. assert.Equal(t, float64(24), resp["user_input"])
  186. // 输出: 120 * 0.8 = 96
  187. assert.Equal(t, float64(96), resp["user_output"])
  188. assert.Equal(t, float64(20), resp["savings_percent"])
  189. assert.Equal(t, "8折", resp["discount"])
  190. }
  191. // TestGetUserPricing_PerCall_WithDiscount 按次计费 + 折扣
  192. func TestGetUserPricing_PerCall_WithDiscount(t *testing.T) {
  193. db := setupPricingTestDB(t)
  194. setPricingCache([]model.Pricing{
  195. {
  196. ModelName: "dall-e-3",
  197. QuotaType: 1,
  198. ModelPrice: 0.04,
  199. EnableGroup: []string{"default", "vip"},
  200. },
  201. })
  202. user := &model.User{Id: 500, Group: "vip", Username: "percallvip", Status: 1}
  203. require.NoError(t, db.Create(user).Error)
  204. withGroupRatio(t, `{"default":1,"vip":0.5}`)
  205. router := setupAuthRouter(500)
  206. w := httptest.NewRecorder()
  207. req, _ := http.NewRequest("GET", "/api/pricing/user/dall-e-3", nil)
  208. router.ServeHTTP(w, req)
  209. var resp map[string]interface{}
  210. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  211. assert.True(t, resp["success"].(bool))
  212. assert.Equal(t, float64(0.04), resp["original_price"])
  213. assert.Equal(t, float64(0.02), resp["user_price"])
  214. assert.Equal(t, float64(50), resp["savings_percent"])
  215. assert.Equal(t, "5折", resp["discount"])
  216. }
  217. // TestFormatDiscount 折扣格式化测试
  218. func TestFormatDiscount(t *testing.T) {
  219. tests := []struct {
  220. ratio float64
  221. expected string
  222. }{
  223. {0.5, "5折"},
  224. {0.8, "8折"},
  225. {0.9, "9折"},
  226. {0.85, "8.5折"},
  227. {0.75, "7.5折"},
  228. {0.95, "9.5折"},
  229. {1.0, ""},
  230. {0.0, "免费"},
  231. }
  232. for _, tt := range tests {
  233. result := formatDiscount(tt.ratio)
  234. assert.Equal(t, tt.expected, result, "ratio=%.2f", tt.ratio)
  235. }
  236. }