new-api/relay/helper/matrix_usage_capability_test.go

package helper

import (
	"testing"

	"github.com/QuantumNous/new-api/constant"
	"github.com/stretchr/testify/require"
)

func TestSupportsMatrixUsageBilling(t *testing.T) {
	require.True(t, SupportsMatrixUsageBilling("seedance-2", constant.ChannelTypeDoubaoVideo, "seedance-2"))
	require.False(t, SupportsMatrixUsageBilling("seedance-2", constant.ChannelTypeOpenAI, "seedance-2"))
	require.False(t, SupportsMatrixUsageBilling("unknown", constant.ChannelTypeDoubaoVideo, "unknown"))
}

func TestSupportsMatrixUsageBilling_DoubaoSeedanceOfficialModelNames(t *testing.T) {
	require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-260128", constant.ChannelTypeDoubaoVideo, "doubao-seedance-2-0-260128"))
	require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-fast-260128", constant.ChannelTypeDoubaoVideo, "doubao-seedance-2-0-fast-260128"))
}

func TestSupportsMatrixUsageBilling_DoubaoAipingChannel(t *testing.T) {
	require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-260128", constant.ChannelTypeDoubaoVideoCompatibleAiping, "doubao-seedance-2-0-260128"))
	require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-fast-260128", constant.ChannelTypeDoubaoVideoCompatibleAiping, "doubao-seedance-2-0-fast-260128"))
}

func TestSupportsMatrixUsageBilling_TianyiYunSeedanceChannel(t *testing.T) {
	require.True(t, SupportsMatrixUsageBilling("cdance2.0-0611", constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-0611"))
	require.True(t, SupportsMatrixUsageBilling("cdance2.0-fast-0611", constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-fast-0611"))
	require.False(t, SupportsMatrixUsageBilling("cdance2.0-0611", constant.ChannelTypeDoubaoVideoCompatibleAiping, "cdance2.0-0611"))
}

func TestResolvePricingModelMapping(t *testing.T) {
	mapped, err := ResolvePricingModelMapping("seedance-2", `{"seedance-2":"doubao-video-seedance"}`)
	require.NoError(t, err)
	require.Equal(t, "doubao-video-seedance", mapped)

	mapped, err = ResolvePricingModelMapping("seedance-2", `{"seedance-2":"a","a":"b"}`)
	require.NoError(t, err)
	require.Equal(t, "b", mapped)

	_, err = ResolvePricingModelMapping("seedance-2", `{"seedance-2":"a","a":"seedance-2"}`)
	require.ErrorContains(t, err, "model_mapping_contains_cycle")
}