easyai-ai-gateway/apps/api/internal/runner/media_parameter_normalization_test.go

package runner

import (
	"testing"

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

func TestParamProcessorTreatsNonRatioValuesAsUnspecifiedForImagesAndVideos(t *testing.T) {
	for _, modelType := range []string{"image_generate", "video_generate"} {
		t.Run(modelType, func(t *testing.T) {
			result := preprocessRequestWithLog(mediaKindForTest(modelType), map[string]any{
				"aspect_ratio": "auto",
				"aspectRatio":  "adaptive",
				"ratio":        "keep_ratio",
			}, store.RuntimeModelCandidate{
				ModelType: modelType,
				Capabilities: map[string]any{
					modelType: map[string]any{
						"aspect_ratio_allowed": []any{"adaptive"},
					},
				},
			})
			for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
				if _, ok := result.Body[key]; ok {
					t.Fatalf("%s should be removed for non-ratio input: %+v", key, result.Body)
				}
			}
			if result.Err != nil {
				t.Fatalf("non-ratio input should use provider default instead of failing: %v", result.Err)
			}
		})
	}
}

func TestParamProcessorNormalizesImageKSizeAndPreservesProviderResolutionChoice(t *testing.T) {
	tests := []struct {
		name       string
		capability map[string]any
		wantSize   string
	}{
		{
			name: "dimension upstream",
			capability: map[string]any{
				"output_resolutions":   []any{"1K", "2K", "4K"},
				"aspect_ratio_allowed": []any{"16:9"},
			},
			wantSize: "2048x1152",
		},
		{
			name: "resolution upstream",
			capability: map[string]any{
				"output_resolutions":   []any{"1K", "2K"},
				"aspect_ratio_allowed": []any{"16:9"},
				"size_param_format":    "resolution",
			},
			wantSize: "2K",
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			result := preprocessRequestWithLog("images.generations", map[string]any{
				"size":         "2k",
				"aspect_ratio": "16 : 9",
			}, store.RuntimeModelCandidate{
				ModelType: "image_generate",
				Capabilities: map[string]any{
					"image_generate": test.capability,
				},
			})
			if result.Err != nil {
				t.Fatalf("preprocess image size: %v", result.Err)
			}
			if result.Body["resolution"] != "2K" || result.Body["size"] != test.wantSize {
				t.Fatalf("unexpected normalized size: %+v", result.Body)
			}
			if result.Body["width"] != 2048 || result.Body["height"] != 1152 || result.Body["aspect_ratio"] != "16:9" {
				t.Fatalf("unexpected normalized dimensions: %+v", result.Body)
			}
		})
	}
}

func TestParamProcessorNormalizesVideo4KSize(t *testing.T) {
	result := preprocessRequestWithLog("videos.generations", map[string]any{
		"size":         "4K",
		"aspect_ratio": "auto",
	}, store.RuntimeModelCandidate{
		ModelType: "video_generate",
		Capabilities: map[string]any{
			"video_generate": map[string]any{
				"output_resolutions":   []any{"2160p"},
				"aspect_ratio_allowed": []any{"16:9"},
			},
		},
	})
	if result.Body["resolution"] != "2160p" {
		t.Fatalf("video 4K size should normalize to 2160p: %+v", result.Body)
	}
	if _, ok := result.Body["aspect_ratio"]; ok {
		t.Fatalf("video auto ratio should stay unspecified: %+v", result.Body)
	}
}

func TestCandidateFilterIgnoresNonRatioValuesAndFiltersRealRatios(t *testing.T) {
	candidates := []store.RuntimeModelCandidate{
		candidateWithAspectRatios("square", "image_generate", "1:1"),
		candidateWithAspectRatios("landscape", "image_generate", "16:9"),
	}

	filtered, summary, err := filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
		"aspect_ratio": "auto",
	}, candidates)
	if err != nil || len(filtered) != 2 || summary != nil {
		t.Fatalf("auto ratio should not participate in candidate matching: filtered=%+v summary=%+v err=%v", filtered, summary, err)
	}

	filtered, summary, err = filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
		"aspect_ratio": "16 : 9",
	}, candidates)
	if err != nil || len(filtered) != 1 || filtered[0].PlatformKey != "landscape" {
		t.Fatalf("real ratio should filter candidates: filtered=%+v summary=%+v err=%v", filtered, summary, err)
	}
}

func TestCandidateFilterNormalizesKSizeAndAllowsUnconfiguredResolution(t *testing.T) {
	candidates := []store.RuntimeModelCandidate{
		candidateWithImageResolutions("one-k", "1K"),
		{
			PlatformID:      "platform-unrestricted",
			PlatformKey:     "unrestricted",
			PlatformModelID: "model-unrestricted",
			ModelName:       "demo-image",
			ModelType:       "image_generate",
			Capabilities: map[string]any{
				"image_generate": map[string]any{},
			},
		},
	}

	filtered, _, err := filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
		"size": "2k",
	}, candidates)
	if err != nil {
		t.Fatalf("2k size should be accepted and normalized for matching: %v", err)
	}
	if len(filtered) != 1 || filtered[0].PlatformKey != "unrestricted" {
		t.Fatalf("unexpected candidates for normalized 2K size: %+v", filtered)
	}
}

func candidateWithAspectRatios(platformKey string, modelType string, ratios ...string) store.RuntimeModelCandidate {
	return store.RuntimeModelCandidate{
		PlatformID:      "platform-" + platformKey,
		PlatformKey:     platformKey,
		PlatformModelID: "model-" + platformKey,
		ModelName:       "demo-media",
		ModelType:       modelType,
		Capabilities: map[string]any{
			modelType: map[string]any{
				"aspect_ratio_allowed": stringsToAny(ratios),
			},
		},
	}
}

func mediaKindForTest(modelType string) string {
	if modelType == "image_generate" {
		return "images.generations"
	}
	return "videos.generations"
}