feat: publish v1.0.0-rc.36-custom.1

controller/billing_option_test.go

@@ -0,0 +1,203 @@
+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.AuditLog{}, &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")
+}