new-api/service/video_asset_channel_admin_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 TestGetUserConcreteTokenGroups(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	for i, group := range []string{"vip", "default", "default", "auto", ""} {
		require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "key-" + string(rune('0'+i)), Group: group}).Error)
	}

	groups, err := model.GetUserConcreteTokenGroups(10)

	require.NoError(t, err)
	assert.Equal(t, []string{"default", "vip"}, groups)
}

func TestGetVideoAssetChannelCandidatesFiltersGroupAndStatus(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 17, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusAutoDisabled)
	createDoubaoAssetChannelForTest(t, db, 18, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "vip", "key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 19, constant.ChannelTypeChinaMobileSeedance, "default", "\n", common.ChannelStatusEnabled)

	candidates, err := GetVideoAssetChannelCandidates("default", VideoAssetFamilySeedance)

	require.NoError(t, err)
	require.Len(t, candidates, 2)
	assert.Equal(t, []int{7, 16}, []int{candidates[0].Id, candidates[1].Id})
}

func TestClearVideoAssetChannelBindingClearsOnlyFamily(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7))
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59))

	require.NoError(t, ClearVideoAssetChannelBinding(10, "default", VideoAssetFamilySeedance))

	seedanceBindings, err := model.GetUserAssetChannelsByTypes(10, VideoAssetChannelTypesForFamily(VideoAssetFamilySeedance), "default")
	require.NoError(t, err)
	assert.Empty(t, seedanceBindings)
	klingBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default")
	require.NoError(t, err)
	assert.NotNil(t, klingBinding)
}