new-api/controller/model_display_pricing_controller_test.go

package controller

import (
	"net/http"
	"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/require"
	"gorm.io/gorm"
)

func setupModelDisplayPricingControllerTest(t *testing.T) *gin.Engine {
	t.Helper()
	gin.SetMode(gin.TestMode)

	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&model.Option{}))
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)

	origDB := model.DB
	origUsingSQLite := common.UsingSQLite
	origRedisEnabled := common.RedisEnabled
	model.DB = db
	common.UsingSQLite = true
	common.RedisEnabled = false
	model.InitOptionMap()
	require.NoError(t, ratio_setting.UpdateModelDisplayPricingByJSONString(`{}`))

	t.Cleanup(func() {
		model.DB = origDB
		common.UsingSQLite = origUsingSQLite
		common.RedisEnabled = origRedisEnabled
		require.NoError(t, ratio_setting.UpdateModelDisplayPricingByJSONString(`{}`))
		require.NoError(t, sqlDB.Close())
	})

	r := gin.New()
	r.GET("/api/option/model_display_pricing", GetModelDisplayPricingRules)
	r.GET("/api/option/model_display_pricing/*model", GetModelDisplayPricingRule)
	r.PUT("/api/option/model_display_pricing/*model", UpdateModelDisplayPricingRule)
	r.DELETE("/api/option/model_display_pricing/*model", DeleteModelDisplayPricingRule)
	return r
}

func TestModelDisplayPricingControllerCRUDAndWildcard(t *testing.T) {
	r := setupModelDisplayPricingControllerTest(t)
	items := []types.ModelDisplayPricingItem{
		{Specification: "768p-6s", OfficialSupplierTip: "768P 6s", Price: 0.33333, Unit: "second", SortOrder: 1},
	}

	w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_display_pricing/provider/kling-video", items)
	require.Equal(t, http.StatusOK, w.Code)
	require.Len(t, ratio_setting.GetModelDisplayPricing("provider/kling-video"), 1)

	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_display_pricing/provider/kling-video", nil)
	require.Equal(t, http.StatusOK, w.Code)
	require.Contains(t, w.Body.String(), "768p-6s")

	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_display_pricing", nil)
	require.Equal(t, http.StatusOK, w.Code)
	require.Contains(t, w.Body.String(), "provider/kling-video")

	w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_display_pricing/provider/kling-video", nil)
	require.Equal(t, http.StatusOK, w.Code)
	require.Empty(t, ratio_setting.GetModelDisplayPricing("provider/kling-video"))
}

func TestModelDisplayPricingControllerValidationError(t *testing.T) {
	r := setupModelDisplayPricingControllerTest(t)
	items := []types.ModelDisplayPricingItem{{Specification: "", Price: 1, Unit: "second"}}

	w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_display_pricing/kling-video", items)
	require.Equal(t, http.StatusBadRequest, w.Code)
	require.Contains(t, w.Body.String(), "specification is required")
}

func TestModelDisplayPricingOptionMapHydratesCache(t *testing.T) {
	r := setupModelDisplayPricingControllerTest(t)
	require.NotNil(t, r)

	require.NoError(t, model.UpdateOption(ratio_setting.ModelDisplayPricingOptionKey, `{
		"kling-video":[{"specification":"768p","price":1,"unit":"second"}]
	}`))

	items := ratio_setting.GetModelDisplayPricing("kling-video")
	require.Len(t, items, 1)
	require.Equal(t, "768p", items[0].Specification)
}