new-api/controller/pricing_user_test.go

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/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"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.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"])
}

func TestGetUserPricingIncludesDisplayPricingWithoutChangingComputedPrice(t *testing.T) {
	setupPricingTestDB(t)
	modelName := "display-pricing-user-model"
	router := gin.New()
	router.GET("/api/pricing/user/*model", GetUserPricing)

	setPricingCache([]model.Pricing{
		{
			ModelName:   modelName,
			QuotaType:   model.QuotaTypeByCall,
			ModelPrice:  0.5,
			EnableGroup: []string{"default"},
			DisplayPricing: []types.ModelDisplayPricingItem{
				{Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1, DiscountRate: 1},
			},
		},
	})

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/api/pricing/user/"+modelName, nil)
	router.ServeHTTP(w, req)

	var resp map[string]interface{}
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.Equal(t, http.StatusOK, w.Code)
	require.Contains(t, resp, "display_pricing")
	require.Equal(t, float64(0.5), resp["original_price"])
	require.NotEqual(t, float64(0.33333), resp["original_price"])
}


func TestGetPricingIncludesDisplayPricingForModelSquare(t *testing.T) {
	setupPricingTestDB(t)
	modelName := "display-pricing-square-model"
	router := gin.New()
	router.GET("/api/pricing", GetPricing)

	setPricingCache([]model.Pricing{
		{
			ModelName:   modelName,
			QuotaType:   model.QuotaTypeByCall,
			ModelPrice:  0.5,
			EnableGroup: []string{"default"},
			DisplayPricing: []types.ModelDisplayPricingItem{
				{Specification: "768p-6s", Price: 0.33333, Unit: "second", SortOrder: 1, DiscountRate: 1},
			},
		},
	})
	withGroupRatio(t, `{"default":1}`)

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/api/pricing", nil)
	router.ServeHTTP(w, req)

	var resp map[string]interface{}
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	require.Equal(t, http.StatusOK, w.Code)
	require.Equal(t, true, resp["success"])
	data := resp["data"].([]interface{})
	require.Len(t, data, 1)
	item := data[0].(map[string]interface{})
	require.Equal(t, modelName, item["model_name"])
	require.Contains(t, item, "display_pricing")
	displayPricing := item["display_pricing"].([]interface{})
	require.Len(t, displayPricing, 1)
	require.Equal(t, "768p-6s", displayPricing[0].(map[string]interface{})["specification"])
}

// 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_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)
	}
}