new-api/relay/relay_task_test.go

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
}