test(task-pricing): cover model pricing api and remix snapshots

controller/model_pricing_controller_test.go

@@ -0,0 +1,135 @@
+package controller
+
+import (
+	"bytes"
+	"net/http"
+	"net/http/httptest"
+	"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/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func setupModelPricingControllerTest(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.UpdateModelPricingRulesByJSONString(`{}`))
+	require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"video-test":0.25}`))
+
+	t.Cleanup(func() {
+		model.DB = origDB
+		common.UsingSQLite = origUsingSQLite
+		common.RedisEnabled = origRedisEnabled
+		require.NoError(t, ratio_setting.UpdateModelPricingRulesByJSONString(`{}`))
+		require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
+		require.NoError(t, sqlDB.Close())
+	})
+
+	r := gin.New()
+	r.GET("/api/option/model_pricing", GetModelPricingRules)
+	r.GET("/api/option/model_pricing/:model", GetModelPricingRule)
+	r.PUT("/api/option/model_pricing/:model", UpdateModelPricingRule)
+	r.DELETE("/api/option/model_pricing/:model", DeleteModelPricingRule)
+	return r
+}
+
+func TestModelPricingController_CRUD(t *testing.T) {
+	r := setupModelPricingControllerTest(t)
+	cfg := validControllerPricingConfig()
+
+	w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/video-test", cfg)
+	require.Equal(t, http.StatusOK, w.Code)
+	assert.True(t, responseSuccess(t, w.Body.Bytes()))
+
+	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/video-test", nil)
+	require.Equal(t, http.StatusOK, w.Code)
+	assert.True(t, responseSuccess(t, w.Body.Bytes()))
+	assert.Contains(t, w.Body.String(), `"billing_unit":"per_call"`)
+
+	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing", nil)
+	require.Equal(t, http.StatusOK, w.Code)
+	assert.True(t, responseSuccess(t, w.Body.Bytes()))
+	assert.Contains(t, w.Body.String(), "video-test")
+
+	w = performJSONRequest(t, r, http.MethodDelete, "/api/option/model_pricing/video-test", nil)
+	require.Equal(t, http.StatusOK, w.Code)
+	assert.True(t, responseSuccess(t, w.Body.Bytes()))
+
+	w = performJSONRequest(t, r, http.MethodGet, "/api/option/model_pricing/video-test", nil)
+	assert.Equal(t, http.StatusNotFound, w.Code)
+	legacyPrice, ok := ratio_setting.GetModelPrice("video-test", false)
+	require.True(t, ok)
+	assert.Equal(t, 0.25, legacyPrice)
+}
+
+func TestModelPricingController_PutRejectsBadConfig(t *testing.T) {
+	r := setupModelPricingControllerTest(t)
+	cfg := validControllerPricingConfig()
+	cfg.Table = nil
+
+	w := performJSONRequest(t, r, http.MethodPut, "/api/option/model_pricing/video-test", cfg)
+	assert.Equal(t, http.StatusBadRequest, w.Code)
+	assert.Contains(t, w.Body.String(), "table must not be empty")
+}
+
+func performJSONRequest(t *testing.T, r *gin.Engine, method, path string, body any) *httptest.ResponseRecorder {
+	t.Helper()
+	var payload []byte
+	if body != nil {
+		var err error
+		payload, err = common.Marshal(body)
+		require.NoError(t, err)
+	}
+	req := httptest.NewRequest(method, path, bytes.NewReader(payload))
+	req.Header.Set("Content-Type", "application/json")
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+	return w
+}
+
+func responseSuccess(t *testing.T, payload []byte) bool {
+	t.Helper()
+	var body struct {
+		Success bool `json:"success"`
+	}
+	require.NoError(t, common.Unmarshal(payload, &body))
+	return body.Success
+}
+
+func validControllerPricingConfig() types.PricingConfig {
+	return types.PricingConfig{
+		SchemaVersion:      1,
+		Scope:              types.PricingScopeModel,
+		BillingUnit:        types.BillingUnitPerCall,
+		PreconsumeStrategy: types.PreconsumeStrategyExact,
+		Dimensions: []types.PricingDimension{
+			{Key: "resolution", Source: "request.resolution", Type: "string"},
+		},
+		Table: []types.PricingRow{
+			{"resolution": "720P", "price": 0.1, "source": types.PricingRowSourceManual},
+		},
+		Fallback: types.PricingFallback{Strategy: types.PricingFallbackReject},
+	}
+}

relay/relay_task_test.go

@@ -0,0 +1,148 @@
+package relay
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/model"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/types"
+	"github.com/gin-gonic/gin"
+	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func setupRelayTaskTestDB(t *testing.T) {
+	t.Helper()
+
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	require.NoError(t, err)
+	require.NoError(t, db.AutoMigrate(&model.Task{}, &model.Channel{}))
+	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
+
+	t.Cleanup(func() {
+		model.DB = origDB
+		common.UsingSQLite = origUsingSQLite
+		common.RedisEnabled = origRedisEnabled
+		require.NoError(t, sqlDB.Close())
+	})
+}
+
+func TestResolveOriginTask_RestoresMatrixOriginPricing(t *testing.T) {
+	setupRelayTaskTestDB(t)
+	require.NoError(t, model.DB.Create(&model.Channel{
+		Id:     11,
+		Type:   constant.ChannelTypeMiniMax,
+		Key:    "test-key",
+		Status: common.ChannelStatusEnabled,
+		Name:   "hailuo",
+		Group:  "default",
+		Models: "matrix-video",
+	}).Error)
+	require.NoError(t, model.DB.Create(&model.Task{
+		TaskID:    "task_origin",
+		UserId:    7,
+		ChannelId: 11,
+		Properties: model.Properties{
+			OriginModelName: "matrix-video",
+		},
+		PrivateData: model.TaskPrivateData{
+			BillingContext: &model.TaskBillingContext{
+				BillingMode:       types.BillingModeMatrix,
+				BillingUnit:       types.BillingUnitPer1MTokens,
+				TokenUnitPriceUSD: 0.5,
+				PricingSnapshot: map[string]any{
+					"resolution":   "720P",
+					"price":        0.5,
+					"billing_unit": types.BillingUnitPer1MTokens,
+					"pricing_mode": string(types.PricingModeMatched),
+				},
+				PerCallBilling: false,
+			},
+		},
+	}).Error)
+
+	c := newRelayTaskRemixContext(t, "task_origin")
+	info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
+
+	taskErr := ResolveOriginTask(c, info)
+	require.Nil(t, taskErr)
+	require.NotNil(t, info.OriginPricing)
+	require.NotNil(t, info.LockedChannel)
+	assert.Equal(t, constant.TaskActionRemix, info.Action)
+	assert.Equal(t, "matrix-video", info.OriginModelName)
+	assert.Equal(t, types.BillingModeMatrix, info.OriginPricing.BillingMode)
+	assert.Equal(t, types.BillingUnitPer1MTokens, info.OriginPricing.BillingUnit)
+	assert.Equal(t, 0.5, info.OriginPricing.TokenUnitPriceUSD)
+	assert.Equal(t, "720P", info.OriginPricing.Snapshot["resolution"])
+	lockedChannel, ok := info.LockedChannel.(*model.Channel)
+	require.True(t, ok)
+	assert.Equal(t, 11, lockedChannel.Id)
+}
+
+func TestResolveOriginTask_RestoresLegacyOriginPricing(t *testing.T) {
+	setupRelayTaskTestDB(t)
+	require.NoError(t, model.DB.Create(&model.Channel{
+		Id:     12,
+		Type:   constant.ChannelTypeOpenAI,
+		Key:    "test-key",
+		Status: common.ChannelStatusEnabled,
+		Name:   "sora",
+		Group:  "default",
+		Models: "legacy-video",
+	}).Error)
+	require.NoError(t, model.DB.Create(&model.Task{
+		TaskID:    "task_legacy",
+		UserId:    7,
+		ChannelId: 12,
+		Properties: model.Properties{
+			OriginModelName: "legacy-video",
+		},
+		PrivateData: model.TaskPrivateData{
+			BillingContext: &model.TaskBillingContext{
+				ModelPrice:     0.2,
+				OtherRatios:    map[string]float64{"seconds": 10, "size": 1.666667},
+				PerCallBilling: true,
+			},
+		},
+	}).Error)
+
+	c := newRelayTaskRemixContext(t, "task_legacy")
+	info := &relaycommon.RelayInfo{UserId: 7, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
+
+	taskErr := ResolveOriginTask(c, info)
+	require.Nil(t, taskErr)
+	require.NotNil(t, info.OriginPricing)
+	require.NotNil(t, info.LockedChannel)
+	assert.Equal(t, "legacy-video", info.OriginModelName)
+	assert.Equal(t, 0.2, info.OriginPricing.PriceUSD)
+	assert.Equal(t, map[string]float64{"seconds": 10, "size": 1.666667}, info.OriginPricing.OtherRatios)
+	assert.True(t, info.OriginPricing.PerCallBilling)
+	lockedChannel, ok := info.LockedChannel.(*model.Channel)
+	require.True(t, ok)
+	assert.Equal(t, 12, lockedChannel.Id)
+}
+
+func newRelayTaskRemixContext(t *testing.T, originTaskID string) *gin.Context {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/"+originTaskID+"/remix", nil)
+	c.Params = gin.Params{{Key: "video_id", Value: originTaskID}}
+	return c
+}