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"]) }