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, "/api/v1/multimodal/kling/videos/text2video", route.UpstreamPath)

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

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

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

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

	videoExtend, ok := FindRoute(http.MethodPost, "/v1/videos/video-extend", RouteKindSubmit)
	require.True(t, ok)
	require.Equal(t, ActionVideoExtend, videoExtend.Action)
	require.Equal(t, "/api/v1/multimodal/kling/videos/video-extend", videoExtend.UpstreamPath)

	motionControl, ok := FindRoute(http.MethodPost, "/v1/videos/motion-control", RouteKindSubmit)
	require.True(t, ok)
	require.Equal(t, "", motionControl.BillingModel, "motion-control BillingModel should be empty so client model_name is used")
}

func TestBuildRequestURLUsesAipingKlingPath(t *testing.T) {
	adaptor := &TaskAdaptor{}
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://aiping.cn"},
		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/multimodal/kling/videos/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, "/api/v1/multimodal/kling/videos/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 TestBuildRequestBodyMapsCanonicalModelNameForAiping(t *testing.T) {
	gin.SetMode(gin.TestMode)
	c, _ := gin.CreateTestContext(nil)
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{},
		TaskRelayInfo: &relaycommon.TaskRelayInfo{
			Action: ActionText2Video,
		},
		OriginModelName: "kling-v2-6",
	}
	relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{
		Model: "kling-v2-6",
		Metadata: map[string]any{
			"model_name": "kling-v2-6",
			"prompt":     "prompt",
			"duration":   "5",
		},
	})

	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"])
}

func TestBuildRequestBodyUsesMappedKlingUpstreamModel(t *testing.T) {
	gin.SetMode(gin.TestMode)
	c, _ := gin.CreateTestContext(nil)
	info := &relaycommon.RelayInfo{
		OriginModelName: "public-kling-alias",
		ChannelMeta: &relaycommon.ChannelMeta{
			UpstreamModelName: "kling-v2-6",
			IsModelMapped:     true,
		},
		TaskRelayInfo: &relaycommon.TaskRelayInfo{
			Action: ActionText2Video,
		},
	}
	relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{
		Model: "public-kling-alias",
		Metadata: map[string]any{
			"model_name": "public-kling-alias",
			"prompt":     "prompt",
			"duration":   "5",
		},
	})

	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")
}

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 TestParseTaskResultConvertsFinalUnitDeductionToTokens(t *testing.T) {
	taskInfo, err := (&TaskAdaptor{}).ParseTaskResult([]byte(`{
		"code": 0,
		"message": "success",
		"request_id": "req-1",
		"data": {
			"task_id": "upstream-task",
			"task_status": "succeed",
			"task_result": {
				"videos": [
					{"id": "v1", "url": "https://example.com/video.mp4", "duration": "5"}
				]
			},
			"final_unit_deduction": "1.234567"
		},
		"aiping_id": "internal"
	}`))

	require.NoError(t, err)
	require.Equal(t, 1234567, taskInfo.CompletionTokens)
	require.Equal(t, 1234567, taskInfo.TotalTokens)
}

func TestFinalUnitDeductionTokens(t *testing.T) {
	cases := []struct {
		name  string
		input any
		want  int
	}{
		{"string decimal", "1.234567", 1234567},
		{"float64", float64(1.234567), 1234567},
		{"int", int(2), 2000000},
		{"int64", int64(3), 3000000},
		{"empty string", "", 0},
		{"whitespace string", "  ", 0},
		{"zero float", float64(0), 0},
		{"negative float", float64(-1.5), 0},
		{"nil", nil, 0},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			require.Equal(t, tc.want, finalUnitDeductionTokens(tc.input))
		})
	}
}

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"])
}