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 TestSupportsMatrixUsageBilling_KlingAipingVideoModels(t *testing.T) {
	cases := []struct {
		billingModel  string
		upstreamModel string
	}{
		{"kling-v1", "Kling-V1"},
		{"kling-v1-6", "Kling-V1.6"},
		{"kling-v2-6", "Kling-V2.6"},
		{"kling-v3", "Kling-V3"},
		{"kling-video-o1", "Kling-Video-O1"},
		{"kling-v3-omni", "Kling-V3-Omni"},
	}

	for _, tc := range cases {
		require.True(t, SupportsMatrixUsageBilling(tc.billingModel, constant.ChannelTypeKlingAiping, tc.billingModel), tc.billingModel)
		require.True(t, SupportsMatrixUsageBilling(tc.billingModel, constant.ChannelTypeKlingAiping, tc.upstreamModel), tc.billingModel)
	}
}

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")
}