easyai-ai-gateway/apps/api/internal/httpapi/kling_compat_integration_test.go

package httpapi

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestKlingCompatibilitySimulationFlow(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the Kling compatibility integration flow")
	}
	ctx := context.Background()
	applyMigration(t, ctx, databaseURL)
	db, err := store.Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()

	var upgradedBaseModels int
	if err := db.Pool().QueryRow(ctx, `
SELECT count(*)
FROM base_model_catalog
WHERE provider_key = 'keling'
  AND provider_model_name IN ('kling-video-o1', 'kling-v3-omni')
  AND model_type @> '["video_generate","image_to_video","omni_video"]'::jsonb
  AND capabilities ? 'video_generate'
  AND capabilities ? 'image_to_video'
  AND capabilities ? 'omni_video'
  AND metadata->'rawModel'->'types' @> '["video_generate","image_to_video","omni_video"]'::jsonb`).Scan(&upgradedBaseModels); err != nil {
		t.Fatalf("read upgraded Kling Omni base model capabilities: %v", err)
	}
	if upgradedBaseModels != 2 {
		t.Fatalf("expected both Kling Omni base models to expose base video capabilities, got %d", upgradedBaseModels)
	}

	serverCtx, cancelServer := context.WithCancel(ctx)
	defer cancelServer()
	server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
		AppEnv:            "test",
		HTTPAddr:          ":0",
		DatabaseURL:       databaseURL,
		IdentityMode:      "hybrid",
		JWTSecret:         "test-secret",
		CORSAllowedOrigin: "*",
	}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
	defer server.Close()

	suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
	username := "kling_compat_" + suffix
	password := "password123"
	var registerResponse struct {
		AccessToken string `json:"accessToken"`
	}
	doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
		"username": username,
		"email":    username + "@example.com",
		"password": password,
	}, http.StatusCreated, &registerResponse)
	var apiKeyResponse struct {
		Secret string `json:"secret"`
	}
	doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", registerResponse.AccessToken, map[string]any{
		"name": "Kling compatibility key",
	}, http.StatusCreated, &apiKeyResponse)
	if _, err := db.Pool().Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
		t.Fatalf("promote compatibility user: %v", err)
	}
	var loginResponse struct {
		AccessToken string `json:"accessToken"`
	}
	doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
		"account": username, "password": password,
	}, http.StatusOK, &loginResponse)

	var platform struct {
		ID string `json:"id"`
	}
	doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms", loginResponse.AccessToken, map[string]any{
		"provider":    "keling",
		"platformKey": "kling-compat-" + suffix,
		"name":        "Kling Compatibility Simulation",
		"baseUrl":     "https://api-beijing.klingai.com/v1",
		"authType":    "AccessKey-SecretKey",
		"credentials": map[string]any{"accessKey": "test-ak", "secretKey": "test-sk"},
	}, http.StatusCreated, &platform)
	for _, model := range []string{klingO1Model, klingV3OmniModel} {
		doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
			"canonicalModelKey": "keling:" + model,
			"modelName":         model,
			"providerModelName": model,
			"modelAlias":        model,
			"modelType":         []string{"omni_video"},
			"displayName":       model,
		}, http.StatusCreated, nil)
	}
	var upgradedPlatformModels int
	if err := db.Pool().QueryRow(ctx, `
SELECT count(*)
FROM platform_models
WHERE platform_id = $1::uuid
  AND model_type @> '["video_generate","image_to_video","omni_video"]'::jsonb
  AND capabilities ? 'video_generate'
  AND capabilities ? 'image_to_video'
  AND capabilities ? 'omni_video'`, platform.ID).Scan(&upgradedPlatformModels); err != nil {
		t.Fatalf("read upgraded Kling Omni platform model capabilities: %v", err)
	}
	if upgradedPlatformModels != 2 {
		t.Fatalf("expected both Kling Omni platform models to expose base video capabilities, got %d", upgradedPlatformModels)
	}

	assertGenericVideoGeneration := func(name string, model string, image string, expectedModelType string) {
		t.Helper()
		t.Run(name, func(t *testing.T) {
			request := map[string]any{
				"model":                model,
				"prompt":               "通用视频接口模拟任务",
				"duration":             5,
				"resolution":           "720p",
				"runMode":              "simulation",
				"simulation":           true,
				"simulationDurationMs": 5,
			}
			if image != "" {
				request["image"] = image
			}
			var response struct {
				Task struct {
					Status        string `json:"status"`
					ModelType     string `json:"modelType"`
					ResolvedModel string `json:"resolvedModel"`
				} `json:"task"`
			}
			doJSON(t, server.URL, http.MethodPost, "/api/v1/videos/generations", apiKeyResponse.Secret, request, http.StatusAccepted, &response)
			resolvedModel, resolved := klingV2ProviderModel(response.Task.ResolvedModel)
			if response.Task.Status != "succeeded" || response.Task.ModelType != expectedModelType || !resolved || resolvedModel != model {
				t.Fatalf("generic video request without modelType should use inferred capability: %+v", response.Task)
			}
		})
	}

	for _, model := range []string{klingO1Model, klingV3OmniModel} {
		assertGenericVideoGeneration(model+"-text-to-video", model, "", "video_generate")
		assertGenericVideoGeneration(model+"-image-to-video", model, "https://example.com/first.png", "image_to_video")
	}

	createV1 := func(model string, duration int, externalID string) string {
		t.Helper()
		var response map[string]any
		doJSON(t, server.URL, http.MethodPost, "/kling/v1/videos/omni-video", apiKeyResponse.Secret, map[string]any{
			"model_name":           model,
			"prompt":               "兼容接口模拟任务",
			"duration":             duration,
			"mode":                 "std",
			"external_task_id":     externalID,
			"runMode":              "simulation",
			"simulation":           true,
			"simulationDurationMs": 5,
		}, http.StatusOK, &response)
		if response["code"] != float64(0) {
			t.Fatalf("unexpected V1 response: %#v", response)
		}
		data, _ := response["data"].(map[string]any)
		taskID, _ := data["task_id"].(string)
		if taskID == "" {
			t.Fatalf("V1 response missing task id: %#v", response)
		}
		return taskID
	}

	o1TaskID := createV1(klingO1Model, 5, "compat-o1-"+suffix)
	v3TaskID := createV1(klingV3OmniModel, 15, "compat-v3-"+suffix)
	for _, taskID := range []string{o1TaskID, v3TaskID} {
		waitKlingV1SimulationTask(t, server.URL, apiKeyResponse.Secret, taskID)
	}
	var listResponse map[string]any
	doJSON(t, server.URL, http.MethodGet, "/kling/v1/videos/omni-video?pageNum=1&pageSize=10", apiKeyResponse.Secret, nil, http.StatusOK, &listResponse)
	items, _ := listResponse["data"].([]any)
	if len(items) < 2 {
		t.Fatalf("V1 task list did not return compatibility tasks: %#v", listResponse)
	}

	var duplicateResponse map[string]any
	doJSON(t, server.URL, http.MethodPost, "/kling/v1/videos/omni-video", apiKeyResponse.Secret, map[string]any{
		"model_name": klingO1Model, "prompt": "duplicate", "duration": 5,
		"external_task_id": "compat-o1-" + suffix,
		"runMode":          "simulation", "simulation": true,
	}, http.StatusConflict, &duplicateResponse)
	if duplicateResponse["code"] != float64(1200) || duplicateResponse["error"] != "external_task_id_reused" {
		t.Fatalf("unexpected duplicate external id response: %#v", duplicateResponse)
	}

	var v2Response map[string]any
	doJSON(t, server.URL, http.MethodPost, "/kling/v2/omni-video/kling-v3-omni", apiKeyResponse.Secret, map[string]any{
		"contents":             []any{map[string]any{"type": "prompt", "text": "API 2.0 模拟任务"}},
		"settings":             map[string]any{"duration": 3, "resolution": "720p", "aspect_ratio": "16:9", "audio": "off"},
		"options":              map[string]any{"external_task_id": "compat-v2-" + suffix},
		"runMode":              "simulation",
		"simulation":           true,
		"simulationDurationMs": 5,
	}, http.StatusOK, &v2Response)
	v2Data, _ := v2Response["data"].(map[string]any)
	v2TaskID, _ := v2Data["id"].(string)
	if v2TaskID == "" {
		t.Fatalf("V2 response missing task id: %#v", v2Response)
	}
	waitKlingV2SimulationTask(t, server.URL, apiKeyResponse.Secret, v2TaskID)
}

func waitKlingV1SimulationTask(t *testing.T, baseURL string, apiKey string, taskID string) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		var response map[string]any
		doJSON(t, baseURL, http.MethodGet, "/kling/v1/videos/omni-video/"+taskID, apiKey, nil, http.StatusOK, &response)
		data, _ := response["data"].(map[string]any)
		switch data["task_status"] {
		case "succeed":
			return
		case "failed":
			t.Fatalf("V1 simulation task failed: %#v", response)
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("V1 simulation task %s timed out", taskID)
}

func waitKlingV2SimulationTask(t *testing.T, baseURL string, apiKey string, taskID string) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		var response map[string]any
		doJSON(t, baseURL, http.MethodGet, "/kling/v2/tasks?task_ids="+taskID, apiKey, nil, http.StatusOK, &response)
		items, _ := response["data"].([]any)
		if len(items) == 1 {
			data, _ := items[0].(map[string]any)
			switch data["status"] {
			case "succeeded":
				return
			case "failed":
				t.Fatalf("V2 simulation task failed: %#v", response)
			}
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("V2 simulation task %s timed out", taskID)
}