new-api/service/doubao_asset_channel_test.go

package service

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

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

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

	oldDB := model.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
		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 createDoubaoAssetChannelForTest(t *testing.T, db *gorm.DB, id int, channelType int, group string, key string, status int) {
	t.Helper()

	priority := int64(id)
	weight := uint(10)
	autoBan := 1
	channel := &model.Channel{
		Id:          id,
		Type:        channelType,
		Key:         key,
		Status:      status,
		Name:        fmt.Sprintf("channel-%d", id),
		Group:       group,
		Models:      "seedance-2",
		Priority:    &priority,
		Weight:      &weight,
		AutoBan:     &autoBan,
		CreatedTime: int64(id),
	}
	require.NoError(t, db.Create(channel).Error)
}

func createDoubaoAssetAbilityForTest(t *testing.T, db *gorm.DB, group string, modelName string, channelId int, enabled bool) {
	t.Helper()

	priority := int64(channelId)
	require.NoError(t, db.Create(&model.Ability{
		Group:     group,
		Model:     modelName,
		ChannelId: channelId,
		Enabled:   enabled,
		Priority:  &priority,
		Weight:    10,
	}).Error)
}

func TestResolveDoubaoAssetChannel_PrefersBoundChannelWithKey(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "bound-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1))

	channel, err := ResolveDoubaoAssetChannel(10, "default")

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, 1, channel.Id)
	assert.Equal(t, "bound-key", channel.Key)
}

func TestResolveDoubaoAssetChannel_AutoMatchesAndPersistsWhenUnbound(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled)

	channel, err := ResolveDoubaoAssetChannel(10, "default")

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, 1, channel.Id)

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

func TestResolveDoubaoAssetChannel_RebindsOpenAIMisbindToAiping(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1))

	channel, err := ResolveDoubaoAssetChannel(10, "default")

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, 2, channel.Id)

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

func TestResolveDoubaoAssetChannel_RebindsOnlyCurrentGroupWhenBoundGroupMismatches(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1))
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", 1))

	channel, err := ResolveDoubaoAssetChannel(10, "default")

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, 2, channel.Id)

	defaultBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default")
	require.NoError(t, err)
	require.NotNil(t, defaultBinding)
	assert.Equal(t, 2, defaultBinding.ChannelId)

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

func TestResolveDoubaoAssetChannel_AutoMatchScansPastFirstPage(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled)
	for i := 2; i <= DoubaoAssetChannelPageSize+1; i++ {
		createDoubaoAssetChannelForTest(t, db, i, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", fmt.Sprintf("vip-key-%d", i), common.ChannelStatusEnabled)
	}

	channel, err := ResolveDoubaoAssetChannel(10, "default")

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, 1, channel.Id)
}

func TestResolveDoubaoAssetChannel_ReturnsErrorWhenBindingWriteFails(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled)

	const callbackName = "doubao_asset_test:fail_user_asset_channel_create"
	require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) {
		if _, ok := tx.Statement.Dest.(*model.UserAssetChannel); ok {
			tx.AddError(assert.AnError)
		}
	}))

	channel, err := ResolveDoubaoAssetChannel(10, "default")

	require.Error(t, err)
	assert.Nil(t, channel)
}

func TestHasDoubaoVideoChannelForGroupModel(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 3, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled)
	createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 1, true)
	createDoubaoAssetAbilityForTest(t, db, "vip", "seedance-2", 2, true)
	createDoubaoAssetAbilityForTest(t, db, "default", "disabled-model", 1, false)
	createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 3, true)

	assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "seedance-2"))
	assert.False(t, HasDoubaoVideoChannelForGroupModel("default", "disabled-model"))
	assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2"))
}

func TestAutoMatchDoubaoAssetChannel_IgnoresLegacyDoubaoVideoChannels(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideo, "default", "legacy-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)

	channel, err := AutoMatchDoubaoAssetChannel("default")

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, 2, channel.Id)
	assert.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleAiping, channel.Type)
}