|
- package controller
-
- import (
- "encoding/json"
- "net/http"
- "net/http/httptest"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "github.com/QuantumNous/new-api/setting/ratio_setting"
- "github.com/glebarez/sqlite"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- // setupPricingTestDB 初始化测试数据库
- func setupPricingTestDB(t *testing.T) *gorm.DB {
- t.Helper()
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := model.DB
- model.DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&model.UserChannelRatio{}, &model.User{}))
-
- t.Cleanup(func() {
- model.DB = origDB
- sqlDB.Close()
- })
- return db
- }
-
- // setupPricingTestRouter 创建无认证路由(未登录场景)
- func setupPricingTestRouter() *gin.Engine {
- gin.SetMode(gin.TestMode)
- r := gin.New()
- r.GET("/api/pricing/user/:model", GetUserPricing)
- return r
- }
-
- // setupAuthRouter 创建带用户 ID 注入的路由(已登录场景)
- func setupAuthRouter(userID int) *gin.Engine {
- gin.SetMode(gin.TestMode)
- r := gin.New()
- r.GET("/api/pricing/user/:model", func(c *gin.Context) {
- c.Set("id", userID)
- c.Next()
- }, GetUserPricing)
- return r
- }
-
- // setPricingCache 直接设置定价缓存用于测试
- func setPricingCache(pricing []model.Pricing) {
- model.SetTestPricing(pricing)
- }
-
- // withGroupRatio 临时设置分组倍率,测试结束后恢复
- func withGroupRatio(t *testing.T, jsonStr string) {
- t.Helper()
- original := ratio_setting.GetGroupRatioCopy()
- ratio_setting.UpdateGroupRatioByJSONString(jsonStr)
- t.Cleanup(func() {
- origJSON, _ := json.Marshal(original)
- ratio_setting.UpdateGroupRatioByJSONString(string(origJSON))
- })
- }
-
- // ---- 测试用例 ----
-
- // TestGetUserPricing_ModelNotFound 模型不存在时应返回错误
- func TestGetUserPricing_ModelNotFound(t *testing.T) {
- setupPricingTestDB(t)
- router := setupPricingTestRouter()
- setPricingCache([]model.Pricing{})
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/nonexistent-model", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.False(t, resp["success"].(bool))
- assert.Contains(t, resp["message"], "未找到")
- }
-
- // TestGetUserPricing_NotLoggedIn 未登录用户应只返回原价
- func TestGetUserPricing_NotLoggedIn(t *testing.T) {
- setupPricingTestDB(t)
- router := setupPricingTestRouter()
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "gpt-4o",
- QuotaType: 0,
- ModelRatio: 15,
- CompletionRatio: 4,
- EnableGroup: []string{"default", "vip"},
- },
- })
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, false, resp["logged_in"])
- assert.Equal(t, "gpt-4o", resp["model_name"])
- assert.Equal(t, float64(0), resp["quota_type"])
-
- // 验证原价: model_ratio * 2 = 15 * 2 = 30
- assert.Equal(t, float64(30), resp["original_input"])
- // 输出原价: model_ratio * completion_ratio * 2 = 15 * 4 * 2 = 120
- assert.Equal(t, float64(120), resp["original_output"])
-
- // 不应有用户价字段
- _, hasUserInput := resp["user_input"]
- assert.False(t, hasUserInput)
- }
-
- // TestGetUserPricing_NotLoggedIn_PerCall 按次计费模型,未登录
- func TestGetUserPricing_NotLoggedIn_PerCall(t *testing.T) {
- setupPricingTestDB(t)
- router := setupPricingTestRouter()
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "dall-e-3",
- QuotaType: 1,
- ModelPrice: 0.04,
- EnableGroup: []string{"default"},
- },
- })
-
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/dall-e-3", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, false, resp["logged_in"])
- assert.Equal(t, float64(0.04), resp["original_price"])
- }
-
- // TestGetUserPricing_LoggedIn_NoDiscount 已登录但无折扣(分组倍率=1,无个人倍率)
- func TestGetUserPricing_LoggedIn_NoDiscount(t *testing.T) {
- db := setupPricingTestDB(t)
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "gpt-4o",
- QuotaType: 0,
- ModelRatio: 15,
- CompletionRatio: 4,
- EnableGroup: []string{"default", "vip"},
- },
- })
-
- // 创建测试用户(default 分组,默认倍率为 1)
- user := &model.User{Id: 100, Group: "default", Username: "testuser", Status: 1}
- require.NoError(t, db.Create(user).Error)
-
- router := setupAuthRouter(100)
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, true, resp["logged_in"])
- // default 分组默认倍率为 1,无折扣
- assert.Equal(t, float64(0), resp["savings_percent"])
- }
-
- // TestGetUserPricing_LoggedIn_GroupDiscount 已登录,有分组折扣
- func TestGetUserPricing_LoggedIn_GroupDiscount(t *testing.T) {
- db := setupPricingTestDB(t)
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "gpt-4o",
- QuotaType: 0,
- ModelRatio: 15,
- CompletionRatio: 4,
- EnableGroup: []string{"default", "vip"},
- },
- })
-
- // 创建 VIP 用户
- user := &model.User{Id: 200, Group: "vip", Username: "vipuser", Status: 1}
- require.NoError(t, db.Create(user).Error)
-
- // 设置 VIP 分组倍率为 0.8
- withGroupRatio(t, `{"default":1,"vip":0.8}`)
-
- router := setupAuthRouter(200)
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, true, resp["logged_in"])
- assert.Equal(t, "vip", resp["group"])
- assert.Equal(t, float64(0.8), resp["group_ratio"])
-
- // 用户价 = 原价 * 0.8
- // 输入: 30 * 0.8 = 24
- assert.Equal(t, float64(24), resp["user_input"])
- // 输出: 120 * 0.8 = 96
- assert.Equal(t, float64(96), resp["user_output"])
-
- assert.Equal(t, float64(20), resp["savings_percent"])
- assert.Equal(t, "8折", resp["discount"])
- }
-
- // TestGetUserPricing_LoggedIn_UserChannelRatio 已登录,有用户渠道倍率
- func TestGetUserPricing_LoggedIn_UserChannelRatio(t *testing.T) {
- db := setupPricingTestDB(t)
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "gpt-4o",
- QuotaType: 0,
- ModelRatio: 15,
- CompletionRatio: 4,
- EnableGroup: []string{"default"},
- },
- })
-
- // 创建用户(default 分组,倍率为1)
- user := &model.User{Id: 300, Group: "default", Username: "specialuser", Status: 1}
- require.NoError(t, db.Create(user).Error)
-
- // 插入用户渠道倍率
- ucr := &model.UserChannelRatio{
- UserId: 300,
- ModelName: "gpt-4o",
- ChannelId: 1,
- Ratio: 0.9,
- }
- require.NoError(t, ucr.Insert())
-
- router := setupAuthRouter(300)
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, true, resp["logged_in"])
- // group_ratio=1 * user_channel_ratio=0.9 = 0.9
- assert.Equal(t, float64(0.9), resp["user_channel_ratio"])
- assert.Equal(t, float64(10), resp["savings_percent"])
- // 输入用户价: 30 * 0.9 = 27
- assert.Equal(t, float64(27), resp["user_input"])
- }
-
- // TestGetUserPricing_LoggedIn_BothDiscounts 分组倍率 + 用户渠道倍率叠加
- func TestGetUserPricing_LoggedIn_BothDiscounts(t *testing.T) {
- db := setupPricingTestDB(t)
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "gpt-4o",
- QuotaType: 0,
- ModelRatio: 15,
- CompletionRatio: 4,
- EnableGroup: []string{"default", "vip"},
- },
- })
-
- user := &model.User{Id: 400, Group: "vip", Username: "bothdiscount", Status: 1}
- require.NoError(t, db.Create(user).Error)
-
- // VIP 分组倍率 0.8
- withGroupRatio(t, `{"default":1,"vip":0.8}`)
-
- // 用户渠道倍率 0.9
- ucr := &model.UserChannelRatio{
- UserId: 400,
- ModelName: "gpt-4o",
- ChannelId: 1,
- Ratio: 0.9,
- }
- require.NoError(t, ucr.Insert())
-
- router := setupAuthRouter(400)
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, true, resp["logged_in"])
- // total = 0.8 * 0.9 = 0.72, savings = 28%
- assert.Equal(t, float64(28), resp["savings_percent"])
- // 输入: 30 * 0.72 = 21.6
- assert.Equal(t, float64(21.6), resp["user_input"])
- // 输出: 120 * 0.72 = 86.4
- assert.Equal(t, float64(86.4), resp["user_output"])
- }
-
- // TestGetUserPricing_PerCall_WithDiscount 按次计费 + 折扣
- func TestGetUserPricing_PerCall_WithDiscount(t *testing.T) {
- db := setupPricingTestDB(t)
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "dall-e-3",
- QuotaType: 1,
- ModelPrice: 0.04,
- EnableGroup: []string{"default", "vip"},
- },
- })
-
- user := &model.User{Id: 500, Group: "vip", Username: "percallvip", Status: 1}
- require.NoError(t, db.Create(user).Error)
-
- withGroupRatio(t, `{"default":1,"vip":0.5}`)
-
- router := setupAuthRouter(500)
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/dall-e-3", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- assert.Equal(t, float64(0.04), resp["original_price"])
- assert.Equal(t, float64(0.02), resp["user_price"])
- assert.Equal(t, float64(50), resp["savings_percent"])
- assert.Equal(t, "5折", resp["discount"])
- }
-
- // TestFormatDiscount 折扣格式化测试
- func TestFormatDiscount(t *testing.T) {
- tests := []struct {
- ratio float64
- expected string
- }{
- {0.5, "5折"},
- {0.8, "8折"},
- {0.9, "9折"},
- {0.85, "8.5折"},
- {0.75, "7.5折"},
- {0.95, "9.5折"},
- {1.0, ""},
- {0.0, "免费"},
- }
- for _, tt := range tests {
- result := formatDiscount(tt.ratio)
- assert.Equal(t, tt.expected, result, "ratio=%.2f", tt.ratio)
- }
- }
-
- // TestGetUserPricing_MultipleUserChannelRatios 多个渠道倍率取最低值
- func TestGetUserPricing_MultipleUserChannelRatios(t *testing.T) {
- db := setupPricingTestDB(t)
-
- setPricingCache([]model.Pricing{
- {
- ModelName: "gpt-4o",
- QuotaType: 0,
- ModelRatio: 15,
- CompletionRatio: 4,
- EnableGroup: []string{"default"},
- },
- })
-
- user := &model.User{Id: 600, Group: "default", Username: "multichannel", Status: 1}
- require.NoError(t, db.Create(user).Error)
-
- // 多个渠道倍率,取最低值 0.7
- for i, ratio := range []float64{0.9, 0.7, 0.8} {
- ucr := &model.UserChannelRatio{
- UserId: 600,
- ModelName: "gpt-4o",
- ChannelId: i + 1,
- Ratio: ratio,
- }
- require.NoError(t, ucr.Insert())
- }
-
- router := setupAuthRouter(600)
- w := httptest.NewRecorder()
- req, _ := http.NewRequest("GET", "/api/pricing/user/gpt-4o", nil)
- router.ServeHTTP(w, req)
-
- var resp map[string]interface{}
- require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
- assert.True(t, resp["success"].(bool))
- // 应取最低倍率 0.7
- assert.Equal(t, float64(0.7), resp["user_channel_ratio"])
- assert.Equal(t, float64(30), resp["savings_percent"])
- }
|