new-api/middleware/doubao_asset_binding_test.go

package middleware

import (
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	relayconstant "github.com/QuantumNous/new-api/relay/constant"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func setupDoubaoVideoBindingDB(t *testing.T) *gorm.DB {
	t.Helper()

	oldDB := model.DB
	oldLOGDB := model.LOG_DB
	oldSQLitePath := common.SQLitePath
	oldMemoryCacheEnabled := common.MemoryCacheEnabled
	oldIsMasterNode := common.IsMasterNode
	oldUsingSQLite := common.UsingSQLite
	oldUsingMySQL := common.UsingMySQL
	oldUsingPostgreSQL := common.UsingPostgreSQL
	oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")

	common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"
	common.MemoryCacheEnabled = false
	common.IsMasterNode = false
	common.UsingSQLite = false
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	require.NoError(t, os.Setenv("SQL_DSN", "local"))
	require.NoError(t, model.InitDB())

	model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
	db := model.DB

	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}))

	t.Cleanup(func() {
		_ = sqlDB.Close()
		model.DB = oldDB
		model.LOG_DB = oldLOGDB
		common.SQLitePath = oldSQLitePath
		common.MemoryCacheEnabled = oldMemoryCacheEnabled
		common.IsMasterNode = oldIsMasterNode
		common.UsingSQLite = oldUsingSQLite
		common.UsingMySQL = oldUsingMySQL
		common.UsingPostgreSQL = oldUsingPostgreSQL
		if hadSQLDSN {
			_ = os.Setenv("SQL_DSN", oldSQLDSN)
		} else {
			_ = os.Unsetenv("SQL_DSN")
		}
	})

	return db
}

func doubaoVideoChannelForBindingTest(id int, group string) *model.Channel {
	priority := int64(id)
	weight := uint(10)
	autoBan := 1
	return &model.Channel{
		Id:       id,
		Type:     constant.ChannelTypeDoubaoVideoCompatibleAiping,
		Key:      fmt.Sprintf("key-%d", id),
		Status:   common.ChannelStatusEnabled,
		Name:     fmt.Sprintf("doubao-video-%d", id),
		Group:    group,
		Models:   "seedance-2",
		Priority: &priority,
		Weight:   &weight,
		AutoBan:  &autoBan,
	}
}

func TestPreferBoundDoubaoVideoChannel(t *testing.T) {
	db := setupDoubaoVideoBindingDB(t)
	channel := doubaoVideoChannelForBindingTest(77, "default")
	priority := int64(channel.Id)
	require.NoError(t, db.Create(&model.Ability{
		Group:     "default",
		Model:     "seedance-2",
		ChannelId: channel.Id,
		Enabled:   true,
		Priority:  &priority,
		Weight:    10,
	}).Error)

	assert.True(t, shouldUseBoundDoubaoVideoChannel(nil, channel, "default", "seedance-2"))
}

func TestPersistDoubaoVideoBindingSkipsNonDoubaoWithoutDB(t *testing.T) {
	oldDB := model.DB
	model.DB = nil
	t.Cleanup(func() {
		model.DB = oldDB
	})

	err := persistDoubaoVideoBindingIfNeeded(10, "default", &model.Channel{
		Id:     88,
		Type:   constant.ChannelTypeOpenAI,
		Status: common.ChannelStatusEnabled,
		Key:    "openai-key",
		Group:  "default",
	})

	require.NoError(t, err)
}

func TestPersistDoubaoVideoBindingSkipsBlankAndAutoGroupWithoutDB(t *testing.T) {
	oldDB := model.DB
	model.DB = nil
	t.Cleanup(func() {
		model.DB = oldDB
	})

	require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "", doubaoVideoChannelForBindingTest(89, "")))
	require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "auto", doubaoVideoChannelForBindingTest(90, "auto")))
}

func TestPersistDoubaoVideoBindingWritesBinding(t *testing.T) {
	setupDoubaoVideoBindingDB(t)
	channel := doubaoVideoChannelForBindingTest(99, "vip")

	require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "vip", channel))

	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip")
	require.NoError(t, err)
	require.NotNil(t, binding)
	assert.Equal(t, channel.Id, binding.ChannelId)
}

func TestShouldCheckDoubaoVideoBindingPreferenceSkipsNonVideoRelayMode(t *testing.T) {
	c, _ := gin.CreateTestContext(nil)
	c.Set("relay_mode", relayconstant.RelayModeChatCompletions)

	assert.False(t, shouldCheckDoubaoVideoBindingPreference(c, "default", "seedance-2"))
}

func TestShouldCheckDoubaoVideoBindingPreferenceRequiresVideoSubmitRelayMode(t *testing.T) {
	db := setupDoubaoVideoBindingDB(t)
	channel := doubaoVideoChannelForBindingTest(101, "default")
	require.NoError(t, db.Create(channel).Error)
	require.NoError(t, db.Create(&model.Ability{
		Group:     "default",
		Model:     "seedance-2",
		ChannelId: channel.Id,
		Enabled:   true,
	}).Error)

	c, _ := gin.CreateTestContext(nil)
	c.Set("relay_mode", relayconstant.RelayModeVideoSubmit)

	assert.True(t, shouldCheckDoubaoVideoBindingPreference(c, "default", "seedance-2"))
}

func TestShouldPersistDoubaoVideoBindingRequiresVideoSubmitRelayMode(t *testing.T) {
	c, _ := gin.CreateTestContext(nil)
	c.Set("relay_mode", relayconstant.RelayModeChatCompletions)
	assert.False(t, shouldPersistDoubaoVideoBinding(c))

	c.Set("relay_mode", relayconstant.RelayModeVideoSubmit)
	assert.True(t, shouldPersistDoubaoVideoBinding(c))
}