new-api/relay/channel/task/kling/aiping/adaptor_test.go

package aiping

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

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

func TestRoutesUseOfficialExternalAndAipingUpstreamPaths(t *testing.T) {
	route, ok := FindRoute(http.MethodPost, "/v1/videos/text2video", RouteKindSubmit)
	require.True(t, ok)
	require.Equal(t, ActionText2Video, route.Action)
	require.Equal(t, "/v1/videos/kling/text2video", route.UpstreamPath)

	voice, ok := FindRoute(http.MethodPost, "/v1/kling/general/custom-voices", RouteKindSubmit)
	require.True(t, ok)
	require.Equal(t, "/v1/kling/general/custom-voices", voice.UpstreamPath)

	presets, ok := FindRoute(http.MethodGet, "/v1/kling/general/presets-voices", RouteKindProxy)
	require.True(t, ok)
	require.Equal(t, "/v1/kling/general/presets-voices", presets.UpstreamPath)

	deleteVoices, ok := FindRoute(http.MethodPost, "/v1/kling/general/delete-voices", RouteKindProxy)
	require.True(t, ok)
	require.Equal(t, ActionDeleteVoices, deleteVoices.Action)

	deleteElements, ok := FindRoute(http.MethodPost, "/v1/kling/general/delete-elements", RouteKindProxy)
	require.True(t, ok)
	require.Equal(t, ActionDeleteElements, deleteElements.Action)
}

func TestBuildRequestURLUsesAipingKlingPath(t *testing.T) {
	adaptor := &TaskAdaptor{}
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://aiping.cn/api"},
		TaskRelayInfo: &relaycommon.TaskRelayInfo{
			Action: ActionText2Video,
		},
	}
	adaptor.Init(info)

	url, err := adaptor.BuildRequestURL(info)
	require.NoError(t, err)
	require.Equal(t, "https://aiping.cn/api/v1/videos/kling/text2video", url)
}

func TestFetchTaskUsesSingleTaskPath(t *testing.T) {
	service.InitHttpClient()
	var gotPath string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"code":0,"data":{"task_id":"upstream-task","task_status":"submitted"}}`))
	}))
	defer server.Close()

	resp, err := (&TaskAdaptor{}).FetchTask(server.URL, "key", map[string]any{
		"task_id": "upstream-task",
		"action":  ActionText2Video,
	}, "")
	require.NoError(t, err)
	defer resp.Body.Close()

	require.Equal(t, "/v1/videos/kling/text2video/upstream-task", gotPath)
}

func TestBuildRequestBodyNormalizesOfficialCompatibleFields(t *testing.T) {
	gin.SetMode(gin.TestMode)
	c, _ := gin.CreateTestContext(nil)
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{},
		TaskRelayInfo: &relaycommon.TaskRelayInfo{
			Action: ActionMultiImage2Video,
		},
		OriginModelName: "Kling-V2.6",
	}
	relaycommon.StoreTaskRequest(c, info, ActionMultiImage2Video, relaycommon.TaskSubmitReq{
		Model: "Kling-V2.6",
		Metadata: map[string]any{
			"model":           "Kling-V1.6",
			"model_name":      "Kling-V2.6",
			"seconds":         3,
			"duration":        5,
			"action_control":  map[string]any{"type": "continuous"},
			"uid":             "user-1",
			"create_at":       float64(1750000000000),
			"_standard_model": "Kling-V-2-6",
			"image_list": []any{
				"https://example.com/1.png",
				map[string]any{"image_url": "https://example.com/2.png"},
				map[string]any{"url": "https://example.com/3.png"},
				map[string]any{"base64": "abc"},
			},
		},
	})

	body, err := (&TaskAdaptor{}).BuildRequestBody(c, info)
	require.NoError(t, err)
	data, err := io.ReadAll(body)
	require.NoError(t, err)

	var got map[string]any
	require.NoError(t, common.Unmarshal(data, &got))
	require.Equal(t, "Kling-V2.6", got["model_name"])
	require.NotContains(t, got, "model")
	require.Equal(t, "5", got["duration"])
	require.NotContains(t, got, "seconds")
	require.NotContains(t, got, "action_control")
	require.NotContains(t, got, "uid")
	require.NotContains(t, got, "create_at")
	require.NotContains(t, got, "_standard_model")
	require.Equal(t, []any{
		map[string]any{"image": "https://example.com/1.png"},
		map[string]any{"image": "https://example.com/2.png"},
		map[string]any{"image": "https://example.com/3.png"},
		map[string]any{"image": "abc"},
	}, got["image_list"])
}

func TestValidateRequestRejectsOmniModelOnTextOrImageRoutes(t *testing.T) {
	gin.SetMode(gin.TestMode)
	c, _ := gin.CreateTestContext(nil)
	info := &relaycommon.RelayInfo{
		TaskRelayInfo: &relaycommon.TaskRelayInfo{
			Action: ActionText2Video,
		},
	}
	relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{
		Model: "kling-v3-omni",
		Metadata: map[string]any{
			"model_name": "kling-v3-omni",
		},
	})

	taskErr := (&TaskAdaptor{}).ValidateRequestAndSetAction(c, info)

	require.NotNil(t, taskErr)
	require.Equal(t, http.StatusUnprocessableEntity, taskErr.StatusCode)
	require.Contains(t, taskErr.Message, "/v1/videos/omni-video")
}

func TestSanitizeNativeTaskPayloadUsesPublicTaskIDAndAddsWatermarkURL(t *testing.T) {
	payload := map[string]any{
		"aiping_id": "internal",
		"data": map[string]any{
			"task_id": "899333358055493641",
			"task_result": map[string]any{
				"videos": []any{
					map[string]any{"id": "v1", "url": "https://example.com/video.mp4"},
				},
			},
		},
	}

	sanitizeNativeTaskPayload(payload, "task_public")
	require.NotContains(t, payload, "aiping_id")
	data := payload["data"].(map[string]any)
	require.Equal(t, "task_public", data["task_id"])
	videos := data["task_result"].(map[string]any)["videos"].([]any)
	require.Equal(t, "", videos[0].(map[string]any)["watermark_url"])
}