new-api/service/video_asset_channel_test.go

package service

import (
	"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"
)

func TestGetBoundVideoAssetChannelForModelFindsSeedanceTianyiYunBinding(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled)
	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16))

	ch, err := GetBoundVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance)

	require.NoError(t, err)
	require.NotNil(t, ch)
	assert.Equal(t, 16, ch.Id)
}

func TestGetBoundVideoAssetChannelForModelFindsKlingBindingByFamily(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled)
	createDoubaoAssetAbilityForTest(t, db, "default", "kling-v2-6", 59, true)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59))

	ch, err := GetBoundVideoAssetChannelForModel(10, "default", "kling-v2-6", VideoAssetFamilyKling)

	require.NoError(t, err)
	require.NotNil(t, ch)
	assert.Equal(t, 59, ch.Id)
}

func TestGetBoundVideoAssetChannelForModelRejectsBindingWithoutAbility(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16))

	ch, err := GetBoundVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance)

	require.NoError(t, err)
	assert.Nil(t, ch)
}

func TestBindVideoAssetChannelReplacesOtherFamilyBindings(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled)
	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 7, true)
	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7))
	selected, err := model.CacheGetChannel(16)
	require.NoError(t, err)

	require.NoError(t, BindVideoAssetChannel(10, "default", selected, VideoAssetFamilySeedance))

	oldBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default")
	require.NoError(t, err)
	assert.Nil(t, oldBinding)
	newBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default")
	require.NoError(t, err)
	require.NotNil(t, newBinding)
	assert.Equal(t, 16, newBinding.ChannelId)
}

func TestGetBoundVideoAssetChannelForModelPrefersLatestFamilyBindingForLegacyRows(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled)
	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 7, true)
	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7))
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16))
	require.NoError(t, db.Model(&model.UserAssetChannel{}).
		Where("user_id = ? AND channel_type = ?", 10, constant.ChannelTypeDoubaoVideoCompatibleAiping).
		Update("updated_at", int64(100)).Error)
	require.NoError(t, db.Model(&model.UserAssetChannel{}).
		Where("user_id = ? AND channel_type = ?", 10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun).
		Update("updated_at", int64(200)).Error)

	ch, err := GetBoundVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance)

	require.NoError(t, err)
	require.NotNil(t, ch)
	assert.Equal(t, 16, ch.Id)
}

func TestResolveVideoAssetChannelForModelAutoSelectsAndPersistsFamilyBinding(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled)
	createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true)

	ch, err := ResolveVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance)

	require.NoError(t, err)
	require.NotNil(t, ch)
	assert.Equal(t, 16, ch.Id)
	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default")
	require.NoError(t, err)
	require.NotNil(t, binding)
	assert.Equal(t, 16, binding.ChannelId)
}