package controller import ( "net/http" "net/http/httptest" "strings" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/setting/config" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) { const pluginKey = "billing-save-probe" const modelName = "billing-save-model" source := ` export const meta = { apiVersion: 1, key: "billing-save-probe", name: "Billing Save Probe", version: "1.0.0", author: {name: "Test"}, models: ["billing-save-model"], fetchMode: "per_task", usageSchema: {seconds: {type: "number", unit: "second"}} }; export function buildSubmitRequest() { return {}; } export function parseSubmitResponse() { return {}; } export function buildQueryRequest() { return {}; } export function parseTaskResult() { return {}; } ` _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) }) tests := []struct { name string expression string errorText string }{ { name: "invalid syntax", expression: `tier("base",`, errorText: "expr compile error", }, { name: "undeclared usage key", expression: `tier("base", u("clips") * 0.1)`, errorText: `usage key \"clips\" is not declared`, }, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { expressions, marshalErr := common.Marshal(map[string]string{modelName: testCase.expression}) require.NoError(t, marshalErr) body, marshalErr := common.Marshal(OptionUpdateRequest{ Key: "billing_setting.billing_expr", Value: string(expressions), }) require.NoError(t, marshalErr) recorder := httptest.NewRecorder() context, _ := gin.CreateTestContext(recorder) context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body))) UpdateOption(context) assert.Equal(t, http.StatusOK, recorder.Code) assert.Contains(t, recorder.Body.String(), `"success":false`) assert.Contains(t, recorder.Body.String(), modelName) assert.Contains(t, recorder.Body.String(), testCase.errorText) }) } } func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) { const modelName = "billing-save-model-without-plugin" expressions, err := common.Marshal(map[string]string{ modelName: `u("mode") == "std" ? 1 : 2`, }) require.NoError(t, err) body, err := common.Marshal(OptionUpdateRequest{ Key: "billing_setting.billing_expr", Value: string(expressions), }) require.NoError(t, err) recorder := httptest.NewRecorder() context, _ := gin.CreateTestContext(recorder) context.Request = httptest.NewRequest( http.MethodPut, "/api/option/", strings.NewReader(string(body)), ) UpdateOption(context) assert.Equal(t, http.StatusOK, recorder.Code) assert.Contains(t, recorder.Body.String(), `"success":false`) assert.Contains(t, recorder.Body.String(), modelName) assert.Contains(t, recorder.Body.String(), "mode") assert.Contains(t, recorder.Body.String(), "no task plugin usage schema") } func setupBillingAliasOptionDB(t *testing.T) { t.Helper() previousDB := model.DB previousLogDB := model.LOG_DB previousType := common.MainDatabaseType() previousCache := common.MemoryCacheEnabled previousMap := common.OptionMap previousRedis := common.RedisEnabled database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Option{}, &model.Log{}, &model.User{})) model.DB = database model.LOG_DB = database common.SetMainDatabaseType(common.DatabaseTypeSQLite) common.MemoryCacheEnabled = false common.RedisEnabled = false common.OptionMap = map[string]string{} t.Cleanup(func() { model.DB = previousDB model.LOG_DB = previousLogDB common.SetMainDatabaseType(previousType) common.MemoryCacheEnabled = previousCache common.OptionMap = previousMap common.RedisEnabled = previousRedis model.InitChannelCache() }) } func TestUpdateOptionAliasBillingExprUsesPluginSchema(t *testing.T) { setupBillingAliasOptionDB(t) const pluginKey = "billing-alias-probe" source := ` export const meta = { apiVersion: 1, key: "billing-alias-probe", name: "Billing Alias Probe", version: "1.0.0", author: {name: "Test"}, models: ["declared-model"], fetchMode: "per_task", usageSchema: {seconds: {type: "number", unit: "second"}} }; export function buildSubmitRequest() { return {}; } export function parseSubmitResponse() { return {}; } export function buildQueryRequest() { return {}; } export function parseTaskResult() { return {}; } ` _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) }) mapping := `{"alias-model":"declared-model"}` require.NoError(t, model.DB.Create(&model.Channel{ Id: 1, Type: 54, Key: "key-1", Status: common.ChannelStatusEnabled, Name: "ch-1", Group: "default", Models: "alias-model,declared-model", ModelMapping: &mapping, }).Error) model.InitChannelCache() saved := map[string]string{} require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error { saved[key] = value return nil })) t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) }) putExpr := func(modelName, expression string) *httptest.ResponseRecorder { t.Helper() expressions, marshalErr := common.Marshal(map[string]string{modelName: expression}) require.NoError(t, marshalErr) body, marshalErr := common.Marshal(OptionUpdateRequest{ Key: "billing_setting.billing_expr", Value: string(expressions), }) require.NoError(t, marshalErr) recorder := httptest.NewRecorder() context, _ := gin.CreateTestContext(recorder) context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body))) UpdateOption(context) return recorder } accepted := putExpr("alias-model", `u("seconds")`) assert.Equal(t, http.StatusOK, accepted.Code) assert.Contains(t, accepted.Body.String(), `"success":true`) rejectedKey := putExpr("alias-model", `u("clips")`) assert.Equal(t, http.StatusOK, rejectedKey.Code) assert.Contains(t, rejectedKey.Body.String(), `"success":false`) assert.Contains(t, rejectedKey.Body.String(), `usage key \"clips\" is not declared`) unresolvable := putExpr("unknown-alias-model", `u("seconds")`) assert.Equal(t, http.StatusOK, unresolvable.Code) assert.Contains(t, unresolvable.Body.String(), `"success":false`) assert.Contains(t, unresolvable.Body.String(), "no task plugin usage schema") }