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.Equal(t, constant.ChannelTypeKlingAiping, requiredTaskChannelTypeForRequest(c)) c = newControllerJSONContext(t, "/v1/videos/video-extend", `{}`) c.Request.URL.Path = "/v1/videos/video-extend" require.Equal(t, constant.ChannelTypeKlingAiping, requiredTaskChannelTypeForRequest(c)) c = newControllerJSONContext(t, "/v1/videos/video_123/remix", `{}`) c.Request.URL.Path = "/v1/videos/video_123/remix" require.Zero(t, requiredTaskChannelTypeForRequest(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") }