new-api/relay/channel/task/doubao/adaptor_test.go

package doubao

import (
	"testing"

	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/stretchr/testify/require"
)

func TestParseTaskResult_SuccessWithUsage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	taskInfo, err := adaptor.ParseTaskResult([]byte(`{
		"status":"succeeded",
		"content":{"video_url":"https://example.test/video.mp4"},
		"usage":{"completion_tokens":12,"total_tokens":3456}
	}`))

	require.NoError(t, err)
	require.Equal(t, string(model.TaskStatusSuccess), taskInfo.Status)
	require.Equal(t, "https://example.test/video.mp4", taskInfo.Url)
	require.Equal(t, 12, taskInfo.CompletionTokens)
	require.Equal(t, 3456, taskInfo.TotalTokens)
}

func TestParseTaskResult_SuccessWithoutUsage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	taskInfo, err := adaptor.ParseTaskResult([]byte(`{
		"status":"succeeded",
		"content":{"video_url":"https://example.test/video.mp4"}
	}`))

	require.NoError(t, err)
	require.Equal(t, string(model.TaskStatusSuccess), taskInfo.Status)
	require.Zero(t, taskInfo.TotalTokens)
}

func TestParseTaskResult_NonSuccessStatusesDoNotRequireUsage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	cases := map[string]model.TaskStatus{
		"queued":     model.TaskStatusQueued,
		"processing": model.TaskStatusInProgress,
		"failed":     model.TaskStatusFailure,
	}
	for status, want := range cases {
		taskInfo, err := adaptor.ParseTaskResult([]byte(`{"status":"` + status + `"}`))
		require.NoError(t, err)
		require.Equal(t, string(want), taskInfo.Status)
		require.Zero(t, taskInfo.TotalTokens)
	}
}

func TestConvertToRequestPayload_MetadataSupportsVideoAudioAndTools(t *testing.T) {
	adaptor := &TaskAdaptor{}
	req := &relaycommon.TaskSubmitReq{
		Model:  "doubao-seedance-2-0-260128",
		Prompt: "make a reference video",
		Metadata: map[string]interface{}{
			"content": []interface{}{
				map[string]interface{}{
					"type": "video_url",
					"video_url": map[string]interface{}{
						"url": "https://example.test/input.mp4",
					},
					"role": "reference_video",
				},
				map[string]interface{}{
					"type": "audio_url",
					"audio_url": map[string]interface{}{
						"url": "https://example.test/input.wav",
					},
					"role": "reference_audio",
				},
			},
			"tools": []interface{}{
				map[string]interface{}{"type": "web_search"},
			},
		},
	}

	payload, err := adaptor.convertToRequestPayload(req)

	require.NoError(t, err)
	require.Len(t, payload.Content, 3)
	require.Equal(t, "video_url", payload.Content[0].Type)
	require.Equal(t, "https://example.test/input.mp4", payload.Content[0].VideoURL.URL)
	require.Equal(t, "reference_video", payload.Content[0].Role)
	require.Equal(t, "audio_url", payload.Content[1].Type)
	require.Equal(t, "https://example.test/input.wav", payload.Content[1].AudioURL.URL)
	require.Equal(t, "reference_audio", payload.Content[1].Role)
	require.Len(t, payload.Tools, 1)
	require.Equal(t, "web_search", payload.Tools[0].Type)
}

func TestConvertToRequestPayload_SecondsOverridesDuration(t *testing.T) {
	adaptor := &TaskAdaptor{}
	req := &relaycommon.TaskSubmitReq{
		Model:   "doubao-seedance-1-5-pro-251215",
		Prompt:  "cat yawning",
		Seconds: "8",
		Metadata: map[string]interface{}{
			"duration": 5,
		},
	}

	payload, err := adaptor.convertToRequestPayload(req)

	require.NoError(t, err)
	require.NotNil(t, payload.Duration)
	require.Equal(t, dto.IntValue(8), *payload.Duration)
}

func TestConvertToRequestPayload_PromptAppendedAfterMetadataAndReplacesMetadataText(t *testing.T) {
	adaptor := &TaskAdaptor{}
	req := &relaycommon.TaskSubmitReq{
		Model:  "doubao-seedance-1-5-pro-251215",
		Prompt: "current prompt",
		Images: []string{"https://example.test/first.png"},
		Metadata: map[string]interface{}{
			"content": []interface{}{
				map[string]interface{}{
					"type": "text",
					"text": "metadata prompt",
				},
				map[string]interface{}{
					"type": "image_url",
					"image_url": map[string]interface{}{
						"url": "https://example.test/second.png",
					},
					"role": "last_frame",
				},
			},
		},
	}

	payload, err := adaptor.convertToRequestPayload(req)

	require.NoError(t, err)
	require.Len(t, payload.Content, 2)
	require.Equal(t, "image_url", payload.Content[0].Type)
	require.Equal(t, "https://example.test/second.png", payload.Content[0].ImageURL.URL)
	require.Equal(t, "last_frame", payload.Content[0].Role)
	require.Equal(t, "text", payload.Content[1].Type)
	require.Equal(t, "current prompt", payload.Content[1].Text)
}

func TestParseTaskResult_FailedUsesUpstreamErrorMessage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	taskInfo, err := adaptor.ParseTaskResult([]byte(`{
		"status":"failed",
		"error":{"code":"InvalidParameter","message":"duration is invalid"}
	}`))

	require.NoError(t, err)
	require.Equal(t, string(model.TaskStatusFailure), taskInfo.Status)
	require.Equal(t, "duration is invalid", taskInfo.Reason)
}

func TestConvertToOpenAIVideo_FailedUsesUpstreamError(t *testing.T) {
	adaptor := &TaskAdaptor{}
	task := &model.Task{
		TaskID:    "task_public",
		Status:    model.TaskStatusFailure,
		Progress:  "100%",
		CreatedAt: 100,
		UpdatedAt: 200,
		Properties: model.Properties{
			OriginModelName: "doubao-seedance-1-5-pro-251215",
		},
		Data: []byte(`{
			"status":"failed",
			"error":{"code":"InvalidParameter","message":"duration is invalid"}
		}`),
	}

	data, err := adaptor.ConvertToOpenAIVideo(task)

	require.NoError(t, err)
	require.Contains(t, string(data), `"message":"duration is invalid"`)
	require.Contains(t, string(data), `"code":"InvalidParameter"`)
}

func TestGetModelList_IncludesSeedance20Models(t *testing.T) {
	adaptor := &TaskAdaptor{}
	models := adaptor.GetModelList()

	require.Contains(t, models, "doubao-seedance-2-0-260128")
	require.Contains(t, models, "doubao-seedance-2-0-fast-260128")
}