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