new-api/service/asset_doubao_test.go

package service

import (
	"net/http"
	"net/url"
	"strings"
	"testing"

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

func TestDoubaoVideoAssetAdapter_SupportsAllOperations(t *testing.T) {
	adapter := NewDoubaoVideoAssetAdapter()
	require.Equal(t, "doubao_video_asset", adapter.Name())
	for _, op := range []AssetOperation{
		AssetOperationAssetCreate,
		AssetOperationAssetList,
		AssetOperationAssetGet,
		AssetOperationAssetUpdate,
		AssetOperationAssetDelete,
		AssetOperationAssetGroupCreate,
		AssetOperationAssetGroupList,
		AssetOperationAssetGroupGet,
		AssetOperationAssetGroupUpdate,
		AssetOperationAssetGroupDelete,
	} {
		require.True(t, adapter.Supports(op), "operation %s should be supported", op)
	}
}

func TestBuildDoubaoVideoAssetURL(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	require.NoError(t, db.AutoMigrate(&model.ChannelAssetCredential{}))
	channel := &model.Channel{Type: constant.ChannelTypeDoubaoVideo}
	rawURL, rawQuery, err := buildDoubaoVideoAssetURL(channel, "ListAssets", "")
	require.NoError(t, err)
	require.Equal(t, "Action=ListAssets&Version=2024-01-01", rawQuery)
	parsed, err := url.Parse(rawURL)
	require.NoError(t, err)
	require.Equal(t, "/openApi/portrait", parsed.Path)
	require.Equal(t, "Action=ListAssets&Version=2024-01-01", parsed.RawQuery)
}

func TestBuildDoubaoVideoAssetURL_DedicatedCredentialBaseWins(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	require.NoError(t, db.AutoMigrate(&model.ChannelAssetCredential{}))
	require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{
		ChannelId: 26,
		AccessKey: "AK",
		SecretKey: "SK",
		BaseURL:   "http://14.103.147.238:19220/openApi/portrait",
	}))
	channel := &model.Channel{Id: 26, Type: constant.ChannelTypeDoubaoVideo}

	rawURL, _, err := buildDoubaoVideoAssetURL(channel, "ListAssets", "")
	require.NoError(t, err)
	require.Equal(t, "http://14.103.147.238:19220/openApi/portrait?Action=ListAssets&Version=2024-01-01", rawURL)
}

func TestBuildDoubaoVideoAssetURL_HardcodedOverrideWinsOverCredential(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	require.NoError(t, db.AutoMigrate(&model.ChannelAssetCredential{}))
	require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{
		ChannelId: 26,
		AccessKey: "AK",
		SecretKey: "SK",
		BaseURL:   "http://credential.example/openApi/portrait",
	}))
	old := doubaoVideoAssetEndpointOverride
	doubaoVideoAssetEndpointOverride = "http://override.example:19220/openApi/portrait"
	defer func() { doubaoVideoAssetEndpointOverride = old }()
	channel := &model.Channel{Id: 26, Type: constant.ChannelTypeDoubaoVideo}

	rawURL, _, err := buildDoubaoVideoAssetURL(channel, "ListAssets", "")
	require.NoError(t, err)
	require.Equal(t, "http://override.example:19220/openApi/portrait?Action=ListAssets&Version=2024-01-01", rawURL)
}

func TestBuildDoubaoVideoAssetURL_FallsBackToOfficialDefault(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	require.NoError(t, db.AutoMigrate(&model.ChannelAssetCredential{}))
	channel := &model.Channel{Id: 26, Type: constant.ChannelTypeDoubaoVideo}

	rawURL, _, err := buildDoubaoVideoAssetURL(channel, "ListAssets", "")
	require.NoError(t, err)
	require.Equal(t, "https://ark.cn-beijing.volces.com/openApi/portrait?Action=ListAssets&Version=2024-01-01", rawURL)
}

func TestSignDoubaoVideoAssetRequest_HeaderShape(t *testing.T) {
	req, err := http.NewRequest(http.MethodPost, "http://14.103.147.238:19220/openApi/portrait?Action=ListAssets&Version=2024-01-01", nil)
	require.NoError(t, err)
	signDoubaoVideoAssetRequest(req, "AK-test", "SK-test", []byte(`{"Filter":{"GroupType":"AIGC"}}`), "Action=ListAssets&Version=2024-01-01")

	auth := req.Header.Get("Authorization")
	require.True(t, strings.HasPrefix(auth, "HMAC-SHA256 Credential=AK-test/"), "got %s", auth)
	require.Contains(t, auth, "/cn-beijing/ark/request, SignedHeaders=content-type;host;x-content-sha256;x-date, Signature=")
	require.NotEmpty(t, req.Header.Get("X-Date"))
	require.Len(t, req.Header.Get("X-Date"), 16)
	require.Equal(t, "application/json", req.Header.Get("Content-Type"))
	// body hash header must be the hex sha256 of the raw body
	require.Regexp(t, `^[0-9a-f]{64}$`, req.Header.Get("X-Content-Sha256"))
}

// The signature must change when any signed input changes (body, key, date).
func TestSignDoubaoVideoAssetRequest_DeterministicPerInputs(t *testing.T) {
	build := func(body string) string {
		req, _ := http.NewRequest(http.MethodPost, "http://h/openApi/portrait?Action=A&Version=2024-01-01", nil)
		signDoubaoVideoAssetRequest(req, "AK", "SK", []byte(body), "Action=A&Version=2024-01-01")
		return req.Header.Get("Authorization")
	}
	sig1 := build(`{"a":1}`)
	sig2 := build(`{"a":1}`)
	sig3 := build(`{"a":2}`)
	require.Equal(t, sig1[strings.Index(sig1, "Signature="):], sig2[strings.Index(sig2, "Signature="):])
	require.NotEqual(t, sig1[strings.Index(sig1, "Signature="):], sig3[strings.Index(sig3, "Signature="):])
}

func TestDoubaoVideoAssetAdapter_MissingCredential(t *testing.T) {
	db := setupDoubaoAssetChannelDB(t)
	require.NoError(t, db.AutoMigrate(&model.ChannelAssetCredential{}))
	adapter := NewDoubaoVideoAssetAdapter()
	channel := &model.Channel{Id: 999999, Type: constant.ChannelTypeDoubaoVideo}
	_, assetErr := adapter.DoAssetRequest(t.Context(), channel, AssetRequest{
		Action: AssetActionSpec{Action: "ListAssets", Operation: AssetOperationAssetList},
		RawBody: []byte(`{}`),
	})
	require.NotNil(t, assetErr)
	require.Contains(t, assetErr.Message, "AccessKey and SecretKey are required")
}