new-api/controller/doubao_aiping_video_test.go

package controller

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

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

func TestConvertAipingNativeVideoRequestPreservesNativeContentOrder(t *testing.T) {
	req, err := convertAipingNativeVideoRequest([]byte(`{
		"model":"doubao-seedance-2-0-260128",
		"content":[
			{"type":"text","text":"first prompt"},
			{"type":"image_url","image_url":{"url":"asset://img1"},"role":"reference_image"},
			{"type":"text","text":"second prompt"},
			{"type":"video_url","video_url":{"url":"https://example.test/input.mp4"},"role":"reference_video"}
		],
		"duration":5,
		"resolution":"480p",
		"tools":[{"type":"web_search"}]
	}`))

	require.NoError(t, err)
	require.Equal(t, "doubao-seedance-2-0-260128", req.Model)
	require.Empty(t, req.Prompt)
	require.Equal(t, float64(5), req.Metadata["duration"])
	require.Equal(t, "480p", req.Metadata["resolution"])
	require.NotNil(t, req.Metadata["tools"])
	content, ok := req.Metadata["content"].([]any)
	require.True(t, ok)
	require.Len(t, content, 4)
	require.Equal(t, "text", content[0].(map[string]any)["type"])
	require.Equal(t, "image_url", content[1].(map[string]any)["type"])
	require.Equal(t, "text", content[2].(map[string]any)["type"])
	require.Equal(t, "video_url", content[3].(map[string]any)["type"])
}

func TestConvertAipingNativeVideoRequestPreservesArkDocumentFields(t *testing.T) {
	req, err := convertAipingNativeVideoRequest([]byte(`{
		"model":"doubao-seedance-2-0-260128",
		"content":[
			{"type":"text","text":"开场：海边日落"},
			{"type":"image_url","image_url":{"url":"asset://image-1"},"role":"reference_image"},
			{"type":"video_url","video_url":{"url":"asset://video-1"},"role":"reference_video"},
			{"type":"audio_url","audio_url":{"url":"asset://audio-1"},"role":"reference_audio"},
			{"type":"draft_task","draft_task":{"id":"cgt-draft"}}
		],
		"callback_url":"https://example.test/callback",
		"return_last_frame":true,
		"service_tier":"default",
		"execution_expires_after":3600,
		"generate_audio":false,
		"draft":true,
		"tools":[{"type":"web_search"}],
		"safety_identifier":"user-hash-1",
		"priority":5,
		"resolution":"480p",
		"ratio":"1:1",
		"duration":5,
		"frames":29,
		"seed":11,
		"camera_fixed":false,
		"watermark":true
	}`))

	require.NoError(t, err)
	require.Equal(t, "doubao-seedance-2-0-260128", req.Model)
	require.Empty(t, req.Prompt)
	require.Equal(t, "https://example.test/callback", req.Metadata["callback_url"])
	require.Equal(t, true, req.Metadata["return_last_frame"])
	require.Equal(t, "default", req.Metadata["service_tier"])
	require.Equal(t, float64(3600), req.Metadata["execution_expires_after"])
	require.Equal(t, false, req.Metadata["generate_audio"])
	require.Equal(t, true, req.Metadata["draft"])
	require.Equal(t, "user-hash-1", req.Metadata["safety_identifier"])
	require.Equal(t, float64(5), req.Metadata["priority"])
	require.Equal(t, "480p", req.Metadata["resolution"])
	require.Equal(t, "1:1", req.Metadata["ratio"])
	require.Equal(t, float64(5), req.Metadata["duration"])
	require.Equal(t, float64(29), req.Metadata["frames"])
	require.Equal(t, float64(11), req.Metadata["seed"])
	require.Equal(t, false, req.Metadata["camera_fixed"])
	require.Equal(t, true, req.Metadata["watermark"])

	content, ok := req.Metadata["content"].([]any)
	require.True(t, ok)
	require.Len(t, content, 5)
	require.Equal(t, "audio_url", content[3].(map[string]any)["type"])
	require.Equal(t, "draft_task", content[4].(map[string]any)["type"])
	tools, ok := req.Metadata["tools"].([]any)
	require.True(t, ok)
	require.Equal(t, "web_search", tools[0].(map[string]any)["type"])
}

func TestConvertAipingNativeVideoRequestRejectsMissingModel(t *testing.T) {
	_, err := convertAipingNativeVideoRequest([]byte(`{"content":[{"type":"text","text":"prompt"}]}`))
	require.ErrorContains(t, err, "model")
}

func TestPrepareAipingNativeVideoSubmitStoresConvertedRequest(t *testing.T) {
	c := newControllerJSONContext(t, "/api/v3/contents/generations/tasks", `{
		"model":"doubao-seedance-2-0-260128",
		"content":[{"type":"text","text":"prompt"}]
	}`)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	err := prepareAipingNativeVideoSubmit(c, info)

	require.NoError(t, err)
	stored, err := relaycommon.GetTaskRequest(c)
	require.NoError(t, err)
	require.Equal(t, "doubao-seedance-2-0-260128", stored.Model)
	require.Empty(t, stored.Prompt)
	content, ok := stored.Metadata["content"].([]any)
	require.True(t, ok)
	require.Equal(t, "prompt", content[0].(map[string]any)["text"])
	require.Equal(t, constant.TaskActionGenerate, info.Action)
	require.Equal(t, "doubao-seedance-2-0-260128", info.OriginModelName)
}

func TestBuildAipingNativeFetchResponseRewritesIDAndPreservesUpstreamFields(t *testing.T) {
	task := &model.Task{
		TaskID:     "task_public",
		Status:     model.TaskStatusSuccess,
		CreatedAt:  1781496040,
		UpdatedAt:  1781496278,
		Properties: model.Properties{OriginModelName: "doubao-seedance-2-0-260128"},
		Data: []byte(`{
			"id":"cgt-upstream",
			"aiping_id":"3d2c8c17-36ad-4138-8602-88b3c60e56c6",
			"model":"doubao-seedance-2-0-260128",
			"status":"succeeded",
			"content":{"video_url":"https://example.test/output.mp4"},
			"usage":{"completion_tokens":48400,"total_tokens":48400},
			"created_at":1781496040,
			"updated_at":1781496278,
			"seed":73812,
			"resolution":"480p",
			"ratio":"1:1",
			"duration":5,
			"framespersecond":24,
			"service_tier":"default",
			"execution_expires_after":172800,
			"generate_audio":true,
			"draft":false,
			"priority":0
		}`),
	}

	data, err := buildAipingNativeFetchResponse(task)

	require.NoError(t, err)
	var payload map[string]any
	require.NoError(t, common.Unmarshal(data, &payload))
	require.Equal(t, "task_public", payload["id"])
	require.Equal(t, "succeeded", payload["status"])
	require.Equal(t, "doubao-seedance-2-0-260128", payload["model"])
	require.Equal(t, "https://example.test/output.mp4", payload["content"].(map[string]any)["video_url"])
	require.Equal(t, float64(48400), payload["usage"].(map[string]any)["total_tokens"])
	require.Equal(t, float64(73812), payload["seed"])
	require.Equal(t, "480p", payload["resolution"])
	require.Equal(t, "1:1", payload["ratio"])
	require.Equal(t, float64(5), payload["duration"])
	require.Equal(t, float64(24), payload["framespersecond"])
	require.Equal(t, "default", payload["service_tier"])
	require.Equal(t, float64(172800), payload["execution_expires_after"])
	require.Equal(t, true, payload["generate_audio"])
	require.Equal(t, false, payload["draft"])
	require.Equal(t, float64(0), payload["priority"])
	require.Contains(t, payload, "error")
	require.Nil(t, payload["error"])
	require.NotContains(t, payload, "aiping_id")
}

func TestBuildAipingNativeFetchResponseFallsBackForFailedTask(t *testing.T) {
	task := &model.Task{
		TaskID:     "task_public",
		Status:     model.TaskStatusFailure,
		CreatedAt:  100,
		UpdatedAt:  200,
		Properties: model.Properties{OriginModelName: "doubao-seedance-2-0-260128"},
		Data:       []byte(`{"id":"cgt-upstream","status":"failed","error":{"code":"InvalidParameter","message":"duration is invalid"}}`),
	}

	data, err := buildAipingNativeFetchResponse(task)

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

func newControllerJSONContext(t *testing.T, path string, body string) *gin.Context {
	t.Helper()
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
	c.Request.Header.Set("Content-Type", "application/json")
	return c
}