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