new-api/service/asset_resolver_test.go

package service

import (
	"context"
	"net/http"
	"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"
)

type fakeAssetAdapter struct {
	name       string
	operation  AssetOperation
	statusCode int
}

func (a fakeAssetAdapter) Name() string {
	return a.name
}

func (a fakeAssetAdapter) Supports(op AssetOperation) bool {
	return a.operation == op
}

func (a fakeAssetAdapter) DoAssetRequest(context.Context, *model.Channel, AssetRequest) (*AssetUpstreamResponse, *AssetError) {
	return &AssetUpstreamResponse{StatusCode: a.statusCode, Body: []byte(`{"ok":true}`)}, nil
}

func TestResolveAssetChannelAutoMatchesOnlyRegisteredOperation(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	resetAssetAdapterRegistryForTest(t)
	RegisterAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fakeAssetAdapter{name: "cm", operation: AssetOperationAssetCreate, statusCode: http.StatusOK})
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeChinaMobileSeedance, "default", "cm-key", common.ChannelStatusEnabled)

	ch, adapter, assetErr := ResolveAssetChannelForOperation(10, "default", AssetOperationAssetCreate)

	require.Nil(t, assetErr)
	require.NotNil(t, ch)
	require.NotNil(t, adapter)
	assert.Equal(t, 2, ch.Id)
	assert.Equal(t, "cm", adapter.Name())
}

func TestResolveAssetChannelBoundUnsupportedOperationDoesNotFallback(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	resetAssetAdapterRegistryForTest(t)
	RegisterAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fakeAssetAdapter{name: "cm", operation: AssetOperationAssetCreate})
	RegisterAssetAdapterForTest(constant.ChannelTypeDoubaoVideoCompatibleAiping, fakeAssetAdapter{name: "aiping", operation: AssetOperationAssetGroupCreate})
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeChinaMobileSeedance, "default", "cm-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeChinaMobileSeedance, "default", 1))

	ch, adapter, assetErr := ResolveAssetChannelForOperation(10, "default", AssetOperationAssetGroupCreate)

	assert.Nil(t, ch)
	assert.Nil(t, adapter)
	require.NotNil(t, assetErr)
	assert.Equal(t, AssetErrorOperationNotSupported, assetErr.Type)
}

func TestResolveAssetChannelIgnoresBindingFromUnrelatedFamily(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	resetAssetAdapterRegistryForTest(t)
	RegisterAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fakeAssetAdapter{name: "cm", operation: AssetOperationAssetCreate})
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeKlingAiping, "default", "kling-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeChinaMobileSeedance, "default", "cm-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 1))

	ch, adapter, assetErr := ResolveAssetChannelForOperation(10, "default", AssetOperationAssetCreate)

	require.Nil(t, assetErr)
	require.NotNil(t, ch)
	require.NotNil(t, adapter)
	assert.Equal(t, 2, ch.Id)
}

func TestResolveAssetChannelAutoMatchDoesNotReplaceVideoBinding(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	resetAssetAdapterRegistryForTest(t)
	RegisterAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fakeAssetAdapter{name: "cm", operation: AssetOperationAssetGroupCreate})
	createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeKlingAiping, "default", "kling-key", common.ChannelStatusEnabled)
	createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeChinaMobileSeedance, "default", "cm-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 1))

	ch, _, assetErr := ResolveAssetChannelForOperation(10, "default", AssetOperationAssetGroupCreate)

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

func TestCompatibleAssetAdapterSupportsOnlyDeclaredOperations(t *testing.T) {
	adapter := NewCompatibleAssetAdapter("aiping_asset", []AssetOperation{AssetOperationAssetCreate})

	assert.True(t, adapter.Supports(AssetOperationAssetCreate))
	assert.False(t, adapter.Supports(AssetOperationAssetGroupCreate))
}