new-api/controller/kling_aiping_native_test.go

package controller

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/dto"
	"github.com/QuantumNous/new-api/model"
	klingaiping "github.com/QuantumNous/new-api/relay/channel/task/kling/aiping"
	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 TestKlingAipingSubmitPreparationUsesModelNameBeforeModel(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{
		"model":"Kling-V1.6",
		"model_name":"Kling-V2.6",
		"prompt":"prompt"
	}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
	modelName := resolveKlingAipingModel(payload, route)
	relaycommon.StoreTaskRequest(c, info, route.Action, relaycommon.TaskSubmitReq{
		Model:    modelName,
		Prompt:   stringFromMap(payload, "prompt"),
		Metadata: payload,
	})

	stored, err := relaycommon.GetTaskRequest(c)
	require.NoError(t, err)
	require.Equal(t, "Kling-V2.6", stored.Model)
	require.Equal(t, "prompt", stored.Prompt)
	require.Equal(t, "Kling-V1.6", stored.Metadata["model"])
	require.Equal(t, "Kling-V2.6", stored.Metadata["model_name"])
}

func TestKlingAipingSubmitPreparationDoesNotLockChannel(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/general/custom-voices", `{
		"voice_url":"https://example.com/voice.mp3",
		"voice_name":"voice"
	}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/general/custom-voices", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
	info.OriginModelName = resolveKlingAipingModel(payload, route)
	info.Action = route.Action
	relaycommon.StoreTaskRequest(c, info, route.Action, relaycommon.TaskSubmitReq{
		Model:    info.OriginModelName,
		Metadata: payload,
	})

	require.Equal(t, klingaiping.ModelCustomVoices, info.OriginModelName)
	require.Equal(t, klingaiping.ActionVoicesCreate, info.Action)
	require.Nil(t, info.LockedChannel)
}

func TestConfigureKlingAipingTaskRelayInfoForcesChannelSelection(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{"model_name":"Kling-V2.6","prompt":"prompt"}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	configureKlingAipingTaskRelayInfo(c, info, route, payload)

	require.NotNil(t, info.ChannelMeta)
	require.Zero(t, info.ChannelMeta.ChannelType)
	require.Equal(t, "Kling-V2.6", info.OriginModelName)
	require.Equal(t, klingaiping.ActionText2Video, info.Action)
	require.Nil(t, info.LockedChannel)
}

func TestConfigureKlingAipingTaskRelayInfoStoresDuration(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{"model_name":"Kling-V2.6","prompt":"prompt","duration":5}`)
	payload, err := readJSONPayload(c)
	require.NoError(t, err)
	route, ok := klingaiping.FindRoute(http.MethodPost, "/v1/videos/text2video", klingaiping.RouteKindSubmit)
	require.True(t, ok)
	info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}

	configureKlingAipingTaskRelayInfo(c, info, route, payload)

	stored, err := relaycommon.GetTaskRequest(c)
	require.NoError(t, err)
	require.Equal(t, 5, stored.Duration)
	require.Empty(t, stored.Seconds)
}

func TestRequiredTaskChannelTypeOnlyMatchesKlingAipingNativePaths(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video", `{}`)
	c.Request.URL.Path = "/v1/videos/text2video"
	require.Zero(t, requiredTaskChannelTypeForRequest(c))
	require.ElementsMatch(t,
		service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling),
		allowedTaskChannelTypesForRequest(c),
	)

	c = newControllerJSONContext(t, "/v1/videos/video-extend", `{}`)
	c.Request.URL.Path = "/v1/videos/video-extend"
	require.Zero(t, requiredTaskChannelTypeForRequest(c))
	require.ElementsMatch(t,
		service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling),
		allowedTaskChannelTypesForRequest(c),
	)

	c = newControllerJSONContext(t, "/v1/videos/video_123/remix", `{}`)
	c.Request.URL.Path = "/v1/videos/video_123/remix"
	require.Zero(t, requiredTaskChannelTypeForRequest(c))
	require.Nil(t, allowedTaskChannelTypesForRequest(c))
}

func TestKlingAipingTaskDataObjectUsesPublicTaskIDAndWatermarkURL(t *testing.T) {
	task := &model.Task{
		TaskID:    "task_public",
		Status:    model.TaskStatusSuccess,
		CreatedAt: 100,
		UpdatedAt: 200,
		Data: []byte(`{
			"code":0,
			"aiping_id":"internal",
			"data":{
				"task_id":"899333358055493641",
				"task_status":"succeed",
				"task_result":{"videos":[{"id":"v1","url":"https://example.com/video.mp4","duration":"5.041"}]}
			}
		}`),
	}

	data := taskDataObject(task)
	require.Equal(t, "task_public", data["task_id"])
	taskResult := data["task_result"].(map[string]any)
	videos := taskResult["videos"].([]any)
	require.Equal(t, "", videos[0].(map[string]any)["watermark_url"])
}

func TestDoKlingAipingProxyRequestUsesSelectedContextKey(t *testing.T) {
	service.InitHttpClient()
	var gotAuth string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"code":0,"message":"success","data":[]}`))
	}))
	defer server.Close()

	c := newControllerJSONContext(t, "/v1/general/advanced-presets-elements", ``)
	common.SetContextKey(c, constant.ContextKeyChannelKey, "selected-key")
	route, ok := klingaiping.FindRoute(http.MethodGet, "/v1/general/advanced-presets-elements", klingaiping.RouteKindProxy)
	require.True(t, ok)
	channel := &model.Channel{
		Key:     "raw-channel-key",
		BaseURL: common.GetPointer(server.URL),
	}

	resp, err := doKlingAipingProxyRequest(c, route, channel)
	require.NoError(t, err)
	defer resp.Body.Close()

	require.Equal(t, "Bearer selected-key", gotAuth)
}

func TestCopyProxyResponseNormalizesMsgError(t *testing.T) {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Set(common.RequestIdKey, "req-test")
	resp := &http.Response{
		StatusCode: http.StatusUnauthorized,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(strings.NewReader(`{"code":401,"msg":"unauthorized","data":null}`)),
	}

	copyProxyResponse(c, resp)

	require.Equal(t, http.StatusUnauthorized, w.Code)
	require.Contains(t, w.Body.String(), `"message":"unauthorized"`)
	require.NotContains(t, w.Body.String(), `"msg"`)
	require.Contains(t, w.Body.String(), `"request_id":"req-test"`)
}

func TestCopyProxyResponseNormalizesPlainTextError(t *testing.T) {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Set(common.RequestIdKey, "req-test")
	resp := &http.Response{
		StatusCode: http.StatusMethodNotAllowed,
		Header:     http.Header{"Content-Type": []string{"text/plain"}},
		Body:       io.NopCloser(strings.NewReader("Method Not Allowed")),
	}

	copyProxyResponse(c, resp)

	require.Equal(t, http.StatusMethodNotAllowed, w.Code)
	require.Contains(t, w.Body.String(), `"message":"Method Not Allowed"`)
	require.Contains(t, w.Body.String(), `"request_id":"req-test"`)
}

func TestCopyProxyResponseNormalizesDetailMessageError(t *testing.T) {
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Set(common.RequestIdKey, "req-test")
	resp := &http.Response{
		StatusCode: http.StatusServiceUnavailable,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(strings.NewReader(`{"detail":{"message":"not found","error_type":"not_found"},"aiping_id":"internal"}`)),
	}

	copyProxyResponse(c, resp)

	require.Equal(t, http.StatusServiceUnavailable, w.Code)
	require.Contains(t, w.Body.String(), `"message":"not found"`)
	require.NotContains(t, w.Body.String(), `map[`)
	require.Contains(t, w.Body.String(), `"request_id":"req-test"`)
}

func TestNormalizeKlingAipingTaskErrorMessageExtractsUpstreamJSONMessage(t *testing.T) {
	taskErr := &dto.TaskError{
		Code:       "fail_to_fetch_task",
		Message:    `{"code":400,"message":"ERROR: image download failed","request_id":"upstream"}`,
		StatusCode: http.StatusBadRequest,
	}

	normalizeKlingAipingTaskError(taskErr)

	require.Equal(t, "ERROR: image download failed", taskErr.Message)
}

func TestParseKlingAipingPageBoundaries(t *testing.T) {
	c := newControllerJSONContext(t, "/v1/videos/text2video?pageNum=1001&pageSize=30", `{}`)
	c.Request.URL.RawQuery = "pageNum=1001&pageSize=30"
	_, _, err := parseKlingAipingPage(c)
	require.ErrorContains(t, err, "pageNum")

	c = newControllerJSONContext(t, "/v1/videos/text2video?pageNum=1&pageSize=501", `{}`)
	c.Request.URL.RawQuery = "pageNum=1&pageSize=501"
	_, _, err = parseKlingAipingPage(c)
	require.ErrorContains(t, err, "pageSize")
}