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/stretchr/testify/assert" "github.com/stretchr/testify/require" ) 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 TestUpdateOptionAliasBillingExprUsesPluginSchema(t *testing.T) { database := modelManagementDB(t, "sqlite", "") require.NoError(t, database.AutoMigrate(&model.Log{})) 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"}, image_count: {type: "number", unit: "count"}}, usageProfiles: [{models: ["declared-model"], schema: {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("image_count")`) assert.Equal(t, http.StatusOK, rejectedKey.Code) assert.Contains(t, rejectedKey.Body.String(), `"success":false`) assert.Contains(t, rejectedKey.Body.String(), `usage key \"image_count\" is not declared`) rejectedDeclared := putExpr("declared-model", `u("image_count")`) assert.Contains(t, rejectedDeclared.Body.String(), `"success":false`) assert.Contains(t, rejectedDeclared.Body.String(), `usage key \"image_count\" 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") }