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

package httpapi

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

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

func TestNormalizeKelingOmniRequestMapsOfficialFields(t *testing.T) {
	normalized, err := normalizeKelingOmniRequest(map[string]any{
		"model_name":       "kling-v3-omni",
		"prompt":           "A rainy street with natural ambience",
		"mode":             "pro",
		"sound":            "on",
		"aspect_ratio":     "9:16",
		"duration":         "5",
		"external_task_id": "external-1",
		"watermark_info":   map[string]any{"enabled": true},
	})
	if err != nil {
		t.Fatalf("normalize Kling request: %v", err)
	}
	if normalized["model"] != "kling-3.0-omni" ||
		normalized["modelType"] != "omni_video" ||
		normalized["resolution"] != "1080p" ||
		normalized["aspect_ratio"] != "9:16" ||
		normalized["duration"] != 5 ||
		normalized["audio"] != true ||
		normalized["sound"] != "on" ||
		normalized["watermark"] != true ||
		normalized["external_task_id"] != "external-1" ||
		normalized["_gateway_compatibility"] != kelingOmniCompatibilityMarker {
		t.Fatalf("unexpected normalized request: %+v", normalized)
	}
	content, _ := normalized["content"].([]any)
	if len(content) != 1 {
		t.Fatalf("unexpected content: %+v", normalized["content"])
	}
	text, _ := content[0].(map[string]any)
	if text["type"] != "text" || text["text"] != "A rainy street with natural ambience" {
		t.Fatalf("unexpected text content: %+v", text)
	}
}

func TestNormalizeKelingOmniRequestBuildsMultiShotMedia(t *testing.T) {
	normalized, err := normalizeKelingOmniRequest(map[string]any{
		"model_name":   "kling-3.0-omni",
		"multi_shot":   true,
		"shot_type":    "customize",
		"aspect_ratio": "16:9",
		"duration":     5,
		"multi_prompt": []any{
			map[string]any{"index": 1, "prompt": "First shot", "duration": "2"},
			map[string]any{"index": 2, "prompt": "Second shot", "duration": 3},
		},
		"image_list": []any{
			map[string]any{"image_url": "https://example.com/reference.png"},
		},
		"element_list": []any{
			map[string]any{"element_id": float64(123)},
		},
	})
	if err != nil {
		t.Fatalf("normalize multi-shot request: %v", err)
	}
	if normalized["model"] != "kling-3.0-omni" || normalized["duration"] != 5 || normalized["multi_shot"] != true || normalized["shot_type"] != "customize" {
		t.Fatalf("unexpected multi-shot fields: %+v", normalized)
	}
	content, _ := normalized["content"].([]any)
	if len(content) != 4 {
		t.Fatalf("expected image, element and two shot prompts, got %+v", content)
	}
}

func TestNormalizeKelingOmniRequestRejectsUnsupportedCombinations(t *testing.T) {
	tests := []struct {
		name string
		body map[string]any
	}{
		{
			name: "callback",
			body: map[string]any{"callback_url": "https://example.com/callback"},
		},
		{
			name: "unknown model",
			body: map[string]any{"model_name": "kling-unknown", "prompt": "x", "aspect_ratio": "16:9"},
		},
		{
			name: "o1 duration",
			body: map[string]any{"model_name": "kling-video-o1", "prompt": "x", "aspect_ratio": "16:9", "duration": 11},
		},
		{
			name: "o1 text to video three seconds",
			body: map[string]any{"model_name": "kling-video-o1", "prompt": "x", "aspect_ratio": "16:9", "duration": 3},
		},
		{
			name: "o1 generated audio",
			body: map[string]any{"model_name": "kling-video-o1", "prompt": "x", "sound": "on", "aspect_ratio": "16:9", "duration": 5},
		},
		{
			name: "video sound",
			body: map[string]any{
				"model_name": "kling-v3-omni",
				"prompt":     "edit",
				"sound":      "on",
				"video_list": []any{map[string]any{"video_url": "https://example.com/base.mp4", "refer_type": "base"}},
			},
		},
		{
			name: "first frame ratio",
			body: map[string]any{
				"prompt":       "animate",
				"aspect_ratio": "16:9",
				"image_list":   []any{map[string]any{"image_url": "https://example.com/first.png", "type": "first_frame"}},
			},
		},
	}
	for _, item := range tests {
		t.Run(item.name, func(t *testing.T) {
			_, err := normalizeKelingOmniRequest(item.body)
			if err == nil || err.Code != 1201 && err.Code != 1203 {
				t.Fatalf("expected official compatibility error, got %v", err)
			}
		})
	}
}

func TestNormalizeKelingO1AllowsThreeSecondsWithReferenceImage(t *testing.T) {
	normalized, err := normalizeKelingOmniRequest(map[string]any{
		"model_name":   "kling-video-o1",
		"prompt":       "Use the landscape as a visual reference",
		"aspect_ratio": "16:9",
		"duration":     3,
		"image_list": []any{
			map[string]any{"image_url": "https://placehold.co/1024x1024/png"},
		},
	})
	if err != nil {
		t.Fatalf("reference-image O1 request should allow three seconds: %v", err)
	}
	if normalized["duration"] != 3 {
		t.Fatalf("unexpected duration: %+v", normalized)
	}
}

func TestKelingCompatTaskDataAndOwnership(t *testing.T) {
	createdAt := time.Unix(100, 0)
	task := store.GatewayTask{
		ID:                "task-1",
		Kind:              "videos.generations",
		GatewayUserID:     "user-1",
		Status:            "succeeded",
		FinalChargeAmount: 2.5,
		Request: map[string]any{
			"_gateway_compatibility": kelingOmniCompatibilityMarker,
			"external_task_id":       "external-1",
			"watermark":              true,
		},
		Result: map[string]any{"data": []any{map[string]any{
			"id":            "video-1",
			"url":           "https://example.com/video.mp4",
			"watermark_url": "https://example.com/watermarked.mp4",
			"duration":      "5",
		}}},
		CreatedAt: createdAt,
		UpdatedAt: createdAt.Add(time.Second),
	}
	if !isKelingCompatTask(task) || !kelingCompatTaskOwnedBy(task, &auth.User{GatewayUserID: "user-1"}) {
		t.Fatalf("expected task ownership and compatibility marker")
	}
	if kelingCompatTaskOwnedBy(task, &auth.User{GatewayUserID: "user-2"}) {
		t.Fatalf("cross-user task access must be rejected")
	}
	data := kelingCompatTaskData(task)
	if data["task_status"] != "succeed" || data["final_unit_deduction"] != "2.5" {
		t.Fatalf("unexpected task data: %+v", data)
	}
	result, _ := data["task_result"].(map[string]any)
	videos, _ := result["videos"].([]any)
	video, _ := videos[0].(map[string]any)
	if video["id"] != "video-1" || video["watermark_url"] != "https://example.com/watermarked.mp4" || video["duration"] != "5" {
		t.Fatalf("unexpected compatible video: %+v", video)
	}
}

func TestRequireKelingAPIKeyWritesOfficialAuthEnvelope(t *testing.T) {
	server := &Server{auth: auth.New("secret", "", "")}
	server.auth.LocalAPIKeyVerifier = func(_ context.Context, key string) (*auth.User, error) {
		if key != "sk-gw-valid" {
			return nil, auth.ErrUnauthorized
		}
		return &auth.User{ID: "user-1", GatewayUserID: "user-1", APIKeyID: "key-1", APIKeyScopes: []string{"video"}}, nil
	}
	handler := server.requireKelingAPIKey(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, ok := auth.UserFromContext(r.Context()); !ok {
			t.Fatal("authenticated user is missing")
		}
		w.WriteHeader(http.StatusNoContent)
	}))

	missing := httptest.NewRecorder()
	handler.ServeHTTP(missing, httptest.NewRequest(http.MethodPost, "/v1/videos/omni-video", nil))
	if missing.Code != http.StatusUnauthorized {
		t.Fatalf("missing auth status=%d body=%s", missing.Code, missing.Body.String())
	}
	var missingBody KelingCompatibleEnvelope
	if err := json.Unmarshal(missing.Body.Bytes(), &missingBody); err != nil || missingBody.Code != 1001 || missingBody.RequestID == "" {
		t.Fatalf("unexpected missing auth envelope: %+v err=%v", missingBody, err)
	}

	invalid := httptest.NewRecorder()
	invalidRequest := httptest.NewRequest(http.MethodPost, "/v1/videos/omni-video", nil)
	invalidRequest.Header.Set("Authorization", "Bearer sk-gw-invalid")
	handler.ServeHTTP(invalid, invalidRequest)
	var invalidBody KelingCompatibleEnvelope
	if err := json.Unmarshal(invalid.Body.Bytes(), &invalidBody); err != nil || invalid.Code != http.StatusUnauthorized || invalidBody.Code != 1002 {
		t.Fatalf("unexpected invalid auth envelope: status=%d body=%+v err=%v", invalid.Code, invalidBody, err)
	}

	valid := httptest.NewRecorder()
	validRequest := httptest.NewRequest(http.MethodPost, "/v1/videos/omni-video", nil)
	validRequest.Header.Set("Authorization", "Bearer sk-gw-valid")
	handler.ServeHTTP(valid, validRequest)
	if valid.Code != http.StatusNoContent {
		t.Fatalf("valid API Key status=%d body=%s", valid.Code, valid.Body.String())
	}
}

func TestKelingCompatErrorImplementsError(t *testing.T) {
	err := newKelingCompatError(http.StatusBadRequest, 1201, "invalid")
	if !errors.Is(err, err) || err.Error() != "invalid" {
		t.Fatalf("unexpected error behavior: %v", err)
	}
}

func TestKelingCompatBusinessCodeMapping(t *testing.T) {
	tests := []struct {
		code    string
		message string
		want    int
	}{
		{code: "insufficient_balance", want: 1101},
		{code: "permission_denied", want: 1103},
		{code: "invalid_parameter", want: 1201},
		{code: "no_model_candidate", want: 1203},
		{code: "rate_limit_exceeded", want: 1302},
		{code: "concurrency_limit", want: 1303},
		{code: "network", want: 5001},
		{code: "unknown", want: 5000},
	}
	for _, item := range tests {
		if got := kelingCompatBusinessCode(item.code, item.message); got != item.want {
			t.Fatalf("code=%q message=%q got=%d want=%d", item.code, item.message, got, item.want)
		}
	}
}

func TestKelingCompatCandidatesSupport4K(t *testing.T) {
	candidates := []store.RuntimeModelCandidate{
		{Provider: "keling", Capabilities: map[string]any{"omni_video": map[string]any{"output_resolutions": []any{"720p", "1080p"}}}},
		{Provider: "keling", Capabilities: map[string]any{"omni_video": map[string]any{"output_resolutions": []any{"2160p"}}}},
	}
	if !kelingCompatCandidatesSupport4K(candidates) {
		t.Fatal("expected 2160p Keling capability to enable mode=4k")
	}
	if kelingCompatCandidatesSupport4K(candidates[:1]) {
		t.Fatal("mode=4k must stay disabled without an explicit capability")
	}
}