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

package doubao_aiping

import (
	"io"
	"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 TestBuildRequestURLUsesAipingVideosEndpoint(t *testing.T) {
	adaptor := &TaskAdaptor{}
	adaptor.Init(&relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{
			ChannelType:    constant.ChannelTypeDoubaoVideoCompatibleAiping,
			ChannelBaseUrl: "https://aiping.example.com",
			ApiKey:         "sk-test",
		},
	})

	got, err := adaptor.BuildRequestURL(&relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://aiping.example.com"},
	})

	require.NoError(t, err)
	require.Equal(t, "https://aiping.example.com/api/v1/multimodal/sd/videos/contents/generations/tasks", got)
}

func TestFetchTaskUsesAipingVideosEndpoint(t *testing.T) {
	adaptor := &TaskAdaptor{}
	req, err := adaptor.buildFetchRequest("https://aiping.example.com", "sk-test", map[string]any{"task_id": "cgt-123"})

	require.NoError(t, err)
	require.Equal(t, http.MethodGet, req.Method)
	require.Equal(t, "https://aiping.example.com/api/v1/multimodal/sd/videos/contents/generations/tasks/cgt-123", req.URL.String())
	require.Equal(t, "Bearer sk-test", req.Header.Get("Authorization"))
}

func TestBuildRequestBodyPassesNativeMetadataToUpstream(t *testing.T) {
	adaptor := &TaskAdaptor{}
	c := newTaskRequestContext(t, `{
		"model":"doubao-seedance-2-0-260128",
		"prompt":"current prompt",
		"metadata":{
			"content":[
				{"type":"video_url","video_url":{"url":"https://example.test/input.mp4"},"role":"reference_video"}
			],
			"duration":5
		}
	}`)
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "doubao-seedance-2-0-260128"},
	}

	body, err := adaptor.BuildRequestBody(c, info)
	require.NoError(t, err)
	data, err := io.ReadAll(body)
	require.NoError(t, err)

	require.Contains(t, string(data), `"video_url":{"url":"https://example.test/input.mp4"}`)
	require.Contains(t, string(data), `"duration":5`)
	require.Contains(t, string(data), `"text":"current prompt"`)
}

func TestBuildRequestBodyPreservesNativeInterleavedContentOrder(t *testing.T) {
	adaptor := &TaskAdaptor{}
	c := newTaskRequestContext(t, `{
		"model":"doubao-seedance-2-0-260128",
		"metadata":{
			"content":[
				{"type":"text","text":"开场：海边日落"},
				{"type":"image_url","image_url":{"url":"asset://img1"},"role":"reference_image"},
				{"type":"text","text":"中段：车身环绕特写"},
				{"type":"video_url","video_url":{"url":"asset://vid1"},"role":"reference_video"},
				{"type":"text","text":"结尾：尾灯点亮"}
			],
			"duration":5
		}
	}`)
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "doubao-seedance-2-0-260128"},
	}

	body, err := adaptor.BuildRequestBody(c, info)
	require.NoError(t, err)
	data, err := io.ReadAll(body)
	require.NoError(t, err)
	var payload requestPayload
	require.NoError(t, common.Unmarshal(data, &payload))

	require.Len(t, payload.Content, 5)
	require.Equal(t, "text", payload.Content[0].Type)
	require.Equal(t, "开场：海边日落", payload.Content[0].Text)
	require.Equal(t, "image_url", payload.Content[1].Type)
	require.Equal(t, "asset://img1", payload.Content[1].ImageURL.URL)
	require.Equal(t, "text", payload.Content[2].Type)
	require.Equal(t, "中段：车身环绕特写", payload.Content[2].Text)
	require.Equal(t, "video_url", payload.Content[3].Type)
	require.Equal(t, "asset://vid1", payload.Content[3].VideoURL.URL)
	require.Equal(t, "text", payload.Content[4].Type)
	require.Equal(t, "结尾：尾灯点亮", payload.Content[4].Text)
}

func TestConvertToRequestPayloadPreservesNativeContentTextOrder(t *testing.T) {
	adaptor := &TaskAdaptor{}
	req := &relaycommon.TaskSubmitReq{
		Model:  "doubao-seedance-2-0-260128",
		Prompt: "first prompt\nsecond prompt\nthird prompt",
		Metadata: map[string]interface{}{
			"content": []interface{}{
				map[string]interface{}{"type": "text", "text": "first prompt"},
				map[string]interface{}{"type": "image_url", "image_url": map[string]interface{}{"url": "asset://img1"}, "role": "reference_image"},
				map[string]interface{}{"type": "text", "text": "second prompt"},
				map[string]interface{}{"type": "video_url", "video_url": map[string]interface{}{"url": "asset://vid1"}, "role": "reference_video"},
				map[string]interface{}{"type": "text", "text": "third prompt"},
			},
		},
	}

	body, err := adaptor.convertToRequestPayload(req)

	require.NoError(t, err)
	require.Len(t, body.Content, 5)
	require.Equal(t, "text", body.Content[0].Type)
	require.Equal(t, "first prompt", body.Content[0].Text)
	require.Equal(t, "image_url", body.Content[1].Type)
	require.Equal(t, "text", body.Content[2].Type)
	require.Equal(t, "second prompt", body.Content[2].Text)
	require.Equal(t, "video_url", body.Content[3].Type)
	require.Equal(t, "text", body.Content[4].Type)
	require.Equal(t, "third prompt", body.Content[4].Text)
}

func TestConvertToRequestPayloadPassesArkDocumentFields(t *testing.T) {
	adaptor := &TaskAdaptor{}
	req := &relaycommon.TaskSubmitReq{
		Model:  "doubao-seedance-2-0-260128",
		Prompt: "documented prompt",
		Metadata: map[string]interface{}{
			"content": []interface{}{
				map[string]interface{}{"type": "text", "text": "documented prompt"},
				map[string]interface{}{"type": "image_url", "image_url": map[string]interface{}{"url": "asset://image-1"}, "role": "reference_image"},
				map[string]interface{}{"type": "video_url", "video_url": map[string]interface{}{"url": "asset://video-1"}, "role": "reference_video"},
				map[string]interface{}{"type": "audio_url", "audio_url": map[string]interface{}{"url": "asset://audio-1"}, "role": "reference_audio"},
				map[string]interface{}{"type": "draft_task", "draft_task": map[string]interface{}{"id": "cgt-draft"}},
			},
			"callback_url":            "https://example.test/callback",
			"return_last_frame":       true,
			"service_tier":            "default",
			"execution_expires_after": float64(3600),
			"generate_audio":          false,
			"draft":                   true,
			"tools":                   []interface{}{map[string]interface{}{"type": "web_search"}},
			"safety_identifier":       "user-hash-1",
			"priority":                float64(5),
			"resolution":              "480p",
			"ratio":                   "1:1",
			"duration":                float64(5),
			"frames":                  float64(29),
			"seed":                    float64(11),
			"camera_fixed":            false,
			"watermark":               true,
		},
	}

	body, err := adaptor.convertToRequestPayload(req)
	require.NoError(t, err)
	data, err := common.Marshal(body)
	require.NoError(t, err)
	jsonBody := string(data)

	require.Contains(t, jsonBody, `"image_url":{"url":"asset://image-1"}`)
	require.Contains(t, jsonBody, `"video_url":{"url":"asset://video-1"}`)
	require.Contains(t, jsonBody, `"audio_url":{"url":"asset://audio-1"}`)
	require.Contains(t, jsonBody, `"draft_task":{"id":"cgt-draft"}`)
	require.Contains(t, jsonBody, `"callback_url":"https://example.test/callback"`)
	require.Contains(t, jsonBody, `"return_last_frame":true`)
	require.Contains(t, jsonBody, `"service_tier":"default"`)
	require.Contains(t, jsonBody, `"execution_expires_after":3600`)
	require.Contains(t, jsonBody, `"generate_audio":false`)
	require.Contains(t, jsonBody, `"draft":true`)
	require.Contains(t, jsonBody, `"tools":[{"type":"web_search"}]`)
	require.Contains(t, jsonBody, `"safety_identifier":"user-hash-1"`)
	require.Contains(t, jsonBody, `"priority":5`)
	require.Contains(t, jsonBody, `"resolution":"480p"`)
	require.Contains(t, jsonBody, `"ratio":"1:1"`)
	require.Contains(t, jsonBody, `"duration":5`)
	require.Contains(t, jsonBody, `"frames":29`)
	require.Contains(t, jsonBody, `"seed":11`)
	require.Contains(t, jsonBody, `"camera_fixed":false`)
	require.Contains(t, jsonBody, `"watermark":true`)
}

func TestDoResponseRewritesIDToPublicTaskID(t *testing.T) {
	adaptor := &TaskAdaptor{}
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	info := &relaycommon.RelayInfo{
		OriginModelName: "doubao-seedance-2-0-260128",
		TaskRelayInfo:   &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"},
	}
	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body: io.NopCloser(strings.NewReader(`{
			"id":"cgt-upstream",
			"aiping_id":"internal-aiping-id",
			"status":"queued",
			"created_at":1781496040,
			"model":"doubao-seedance-2-0-260128"
		}`)),
	}

	upstreamID, data, taskErr := adaptor.DoResponse(c, resp, info)

	require.Nil(t, taskErr)
	require.Equal(t, "cgt-upstream", upstreamID)
	require.Contains(t, string(data), `"id":"cgt-upstream"`)
	require.Equal(t, http.StatusOK, w.Code)
	require.Contains(t, w.Body.String(), `"id":"task_public"`)
	require.NotContains(t, w.Body.String(), "aiping_id")
	require.Contains(t, w.Body.String(), `"status":"queued"`)
	require.Contains(t, w.Body.String(), `"model":"doubao-seedance-2-0-260128"`)
}

func TestParseTaskResultTreatsArkTerminalStatusesAsFinished(t *testing.T) {
	adaptor := &TaskAdaptor{}
	cases := []struct {
		name       string
		body       string
		wantReason string
	}{
		{
			name:       "expired",
			body:       `{"id":"cgt-expired","status":"expired","error":{"code":"Expired","message":"task expired"}}`,
			wantReason: "task expired",
		},
		{
			name:       "cancelled",
			body:       `{"id":"cgt-cancelled","status":"cancelled","error":{"code":"Cancelled","message":"task cancelled"}}`,
			wantReason: "task cancelled",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got, err := adaptor.ParseTaskResult([]byte(tc.body))

			require.NoError(t, err)
			require.Equal(t, model.TaskStatusFailure, got.Status)
			require.Equal(t, "100%", got.Progress)
			require.Equal(t, tc.wantReason, got.Reason)
		})
	}
}

func newTaskRequestContext(t *testing.T, body string) *gin.Context {
	t.Helper()
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(body))
	c.Request.Header.Set("Content-Type", "application/json")
	var req relaycommon.TaskSubmitReq
	require.NoError(t, common.Unmarshal([]byte(body), &req))
	relaycommon.StoreTaskRequest(c, &relaycommon.RelayInfo{}, constant.TaskActionGenerate, req)
	return c
}