Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

411 строки
12 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.UserChannelRatio{}, &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_LoggedIn_UserChannelRatio 已登录,有用户渠道倍率
  192. func TestGetUserPricing_LoggedIn_UserChannelRatio(t *testing.T) {
  193. db := setupPricingTestDB(t)
  194. setPricingCache([]model.Pricing{
  195. {
  196. ModelName: "gpt-4o",
  197. QuotaType: 0,
  198. ModelRatio: 15,
  199. CompletionRatio: 4,
  200. EnableGroup: []string{"default"},
  201. },
  202. })
  203. // 创建用户(default 分组,倍率为1)
  204. user := &model.User{Id: 300, Group: "default", Username: "specialuser", Status: 1}
  205. require.NoError(t, db.Create(user).Error)
  206. // 插入用户渠道倍率
  207. ucr := &model.UserChannelRatio{
  208. UserId: 300,
  209. ModelName: "gpt-4o",
  210. ChannelId: 1,
  211. Ratio: 0.9,
  212. }
  213. require.NoError(t, ucr.Insert())
  214. router := setupAuthRouter(300)
  215. w := httptest.NewRecorder()
  216. req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
  217. router.ServeHTTP(w, req)
  218. var resp map[string]interface{}
  219. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  220. assert.True(t, resp["success"].(bool))
  221. assert.Equal(t, true, resp["logged_in"])
  222. // group_ratio=1 * user_channel_ratio=0.9 = 0.9
  223. assert.Equal(t, float64(0.9), resp["user_channel_ratio"])
  224. assert.Equal(t, float64(10), resp["savings_percent"])
  225. // 输入用户价: 30 * 0.9 = 27
  226. assert.Equal(t, float64(27), resp["user_input"])
  227. }
  228. // TestGetUserPricing_LoggedIn_BothDiscounts 分组倍率 + 用户渠道倍率叠加
  229. func TestGetUserPricing_LoggedIn_BothDiscounts(t *testing.T) {
  230. db := setupPricingTestDB(t)
  231. setPricingCache([]model.Pricing{
  232. {
  233. ModelName: "gpt-4o",
  234. QuotaType: 0,
  235. ModelRatio: 15,
  236. CompletionRatio: 4,
  237. EnableGroup: []string{"default", "vip"},
  238. },
  239. })
  240. user := &model.User{Id: 400, Group: "vip", Username: "bothdiscount", Status: 1}
  241. require.NoError(t, db.Create(user).Error)
  242. // VIP 分组倍率 0.8
  243. withGroupRatio(t, `{"default":1,"vip":0.8}`)
  244. // 用户渠道倍率 0.9
  245. ucr := &model.UserChannelRatio{
  246. UserId: 400,
  247. ModelName: "gpt-4o",
  248. ChannelId: 1,
  249. Ratio: 0.9,
  250. }
  251. require.NoError(t, ucr.Insert())
  252. router := setupAuthRouter(400)
  253. w := httptest.NewRecorder()
  254. req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
  255. router.ServeHTTP(w, req)
  256. var resp map[string]interface{}
  257. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  258. assert.True(t, resp["success"].(bool))
  259. assert.Equal(t, true, resp["logged_in"])
  260. // total = 0.8 * 0.9 = 0.72, savings = 28%
  261. assert.Equal(t, float64(28), resp["savings_percent"])
  262. // 输入: 30 * 0.72 = 21.6
  263. assert.Equal(t, float64(21.6), resp["user_input"])
  264. // 输出: 120 * 0.72 = 86.4
  265. assert.Equal(t, float64(86.4), resp["user_output"])
  266. }
  267. // TestGetUserPricing_PerCall_WithDiscount 按次计费 + 折扣
  268. func TestGetUserPricing_PerCall_WithDiscount(t *testing.T) {
  269. db := setupPricingTestDB(t)
  270. setPricingCache([]model.Pricing{
  271. {
  272. ModelName: "dall-e-3",
  273. QuotaType: 1,
  274. ModelPrice: 0.04,
  275. EnableGroup: []string{"default", "vip"},
  276. },
  277. })
  278. user := &model.User{Id: 500, Group: "vip", Username: "percallvip", Status: 1}
  279. require.NoError(t, db.Create(user).Error)
  280. withGroupRatio(t, `{"default":1,"vip":0.5}`)
  281. router := setupAuthRouter(500)
  282. w := httptest.NewRecorder()
  283. req, _ := http.NewRequest("GET", "/api/pricing/user/dall-e-3", nil)
  284. router.ServeHTTP(w, req)
  285. var resp map[string]interface{}
  286. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  287. assert.True(t, resp["success"].(bool))
  288. assert.Equal(t, float64(0.04), resp["original_price"])
  289. assert.Equal(t, float64(0.02), resp["user_price"])
  290. assert.Equal(t, float64(50), resp["savings_percent"])
  291. assert.Equal(t, "5折", resp["discount"])
  292. }
  293. // TestFormatDiscount 折扣格式化测试
  294. func TestFormatDiscount(t *testing.T) {
  295. tests := []struct {
  296. ratio float64
  297. expected string
  298. }{
  299. {0.5, "5折"},
  300. {0.8, "8折"},
  301. {0.9, "9折"},
  302. {0.85, "8.5折"},
  303. {0.75, "7.5折"},
  304. {0.95, "9.5折"},
  305. {1.0, ""},
  306. {0.0, "免费"},
  307. }
  308. for _, tt := range tests {
  309. result := formatDiscount(tt.ratio)
  310. assert.Equal(t, tt.expected, result, "ratio=%.2f", tt.ratio)
  311. }
  312. }
  313. // TestGetUserPricing_MultipleUserChannelRatios 多个渠道倍率取最低值
  314. func TestGetUserPricing_MultipleUserChannelRatios(t *testing.T) {
  315. db := setupPricingTestDB(t)
  316. setPricingCache([]model.Pricing{
  317. {
  318. ModelName: "gpt-4o",
  319. QuotaType: 0,
  320. ModelRatio: 15,
  321. CompletionRatio: 4,
  322. EnableGroup: []string{"default"},
  323. },
  324. })
  325. user := &model.User{Id: 600, Group: "default", Username: "multichannel", Status: 1}
  326. require.NoError(t, db.Create(user).Error)
  327. // 多个渠道倍率,取最低值 0.7
  328. for i, ratio := range []float64{0.9, 0.7, 0.8} {
  329. ucr := &model.UserChannelRatio{
  330. UserId: 600,
  331. ModelName: "gpt-4o",
  332. ChannelId: i + 1,
  333. Ratio: ratio,
  334. }
  335. require.NoError(t, ucr.Insert())
  336. }
  337. router := setupAuthRouter(600)
  338. w := httptest.NewRecorder()
  339. req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
  340. router.ServeHTTP(w, req)
  341. var resp map[string]interface{}
  342. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  343. assert.True(t, resp["success"].(bool))
  344. // 应取最低倍率 0.7
  345. assert.Equal(t, float64(0.7), resp["user_channel_ratio"])
  346. assert.Equal(t, float64(30), resp["savings_percent"])
  347. }