new-api/service/kling_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 TestResolveKlingAssetChannel_PrefersBoundChannel(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeKlingAiping, "default", "bound-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeKlingAiping, "default", "auto-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 1))

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

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

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

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

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

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

func TestResolveKlingAssetChannel_IgnoresDisabledChannel(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeKlingAiping, "default", "disabled-key", common.ChannelStatusAutoDisabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeKlingAiping, "default", "enabled-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 1))

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

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

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

	channel, err := GetBoundKlingAssetChannelForModel(10, "default", "kling-v2-6")

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

func TestGetBoundKlingAssetChannelForModelFindsBoundChannel(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))

	channel, err := GetBoundKlingAssetChannelForModel(10, "default", "kling-v2-6")

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

func TestIsUsableKlingAssetChannel_RejectsOtherChannelType(t *testing.T) {
	ch := &model.Channel{Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "default"}
	assert.False(t, IsUsableKlingAssetChannel(ch, "default"))
}

func TestIsUsableKlingAssetChannel_RejectsGroupMismatch(t *testing.T) {
	ch := &model.Channel{Type: constant.ChannelTypeKlingAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "vip"}
	assert.False(t, IsUsableKlingAssetChannel(ch, "default"))
}