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

package runner

import (
	"bytes"
	"context"
	"encoding/base64"
	"image"
	"image/color"
	"image/jpeg"
	"image/png"
	"strings"
	"testing"
	"time"

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

func seedanceImageConstraintCandidate() store.RuntimeModelCandidate {
	return store.RuntimeModelCandidate{
		ModelType: "image_to_video",
		Capabilities: map[string]any{
			"image_to_video": map[string]any{
				"input_image_resolution_range": map[string]any{
					"min": map[string]any{"long_edge": float64(360), "short_edge": float64(360)},
					"max": map[string]any{"long_edge": float64(1920), "short_edge": float64(1080)},
				},
				"input_image_aspect_ratio_range": []any{float64(0.39), float64(2.5)},
			},
		},
	}
}

func pngImageDataURL(t *testing.T, width int, height int) string {
	t.Helper()
	source := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			source.Set(x, y, color.RGBA{R: 80, G: 120, B: 160, A: 255})
		}
	}
	var payload bytes.Buffer
	if err := png.Encode(&payload, source); err != nil {
		t.Fatal(err)
	}
	return "data:image/png;base64," + base64.StdEncoding.EncodeToString(payload.Bytes())
}

func decodedImageSize(t *testing.T, source string) image.Point {
	t.Helper()
	if !strings.HasPrefix(source, "data:image/jpeg;base64,") {
		t.Fatalf("expected converted JPEG data URL, got prefix %.32q", source)
	}
	payload, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(source, "data:image/jpeg;base64,"))
	if err != nil {
		t.Fatal(err)
	}
	config, err := jpeg.DecodeConfig(bytes.NewReader(payload))
	if err != nil {
		t.Fatal(err)
	}
	return image.Pt(config.Width, config.Height)
}

func jpegImageDataURL(t *testing.T, width int, height int) string {
	t.Helper()
	var payload bytes.Buffer
	if err := jpeg.Encode(&payload, image.NewRGBA(image.Rect(0, 0, width, height)), &jpeg.Options{Quality: 90}); err != nil {
		t.Fatal(err)
	}
	return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(payload.Bytes())
}

func TestNormalizeVideoInputImagesResizesOversizedSeedanceImage(t *testing.T) {
	service := &Service{}
	body := map[string]any{
		"content": []any{
			map[string]any{
				"type":      "image_url",
				"role":      "first_frame",
				"image_url": map[string]any{"url": pngImageDataURL(t, 3840, 2160)},
			},
		},
	}

	normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
	if err != nil {
		t.Fatal(err)
	}
	content := normalized["content"].([]any)
	item := content[0].(map[string]any)
	imageURL := item["image_url"].(map[string]any)["url"].(string)
	if got := decodedImageSize(t, imageURL); got != image.Pt(1920, 1080) {
		t.Fatalf("unexpected converted size: %v", got)
	}
}

func TestNormalizeVideoInputImagesBoundsSixKReferenceMemoryWork(t *testing.T) {
	service := &Service{imageNormalizeSlots: make(chan struct{}, 1)}
	candidate := store.RuntimeModelCandidate{
		ModelType: "omni_video",
		Capabilities: map[string]any{
			"omni_video": map[string]any{
				"input_image_resolution_range": map[string]any{
					"min": map[string]any{"long_edge": float64(300), "short_edge": float64(300)},
					"max": map[string]any{"long_edge": float64(6000), "short_edge": float64(6000)},
				},
				"input_image_aspect_ratio_range": []any{float64(0.4), float64(2.5)},
			},
		},
	}
	body := map[string]any{"first_frame": jpegImageDataURL(t, 6144, 2160)}
	normalized, err := service.normalizeVideoInputImages(context.Background(), body, candidate)
	if err != nil {
		t.Fatal(err)
	}
	if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(6000, 2400) {
		t.Fatalf("unexpected converted size: %v", got)
	}
}

func TestImageNormalizationSlotSerializesHeavyConversions(t *testing.T) {
	service := &Service{imageNormalizeSlots: make(chan struct{}, 1)}
	releaseFirst, err := service.acquireImageNormalizationSlot(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	acquired := make(chan func(), 1)
	go func() {
		release, acquireErr := service.acquireImageNormalizationSlot(context.Background())
		if acquireErr == nil {
			acquired <- release
		}
	}()
	select {
	case release := <-acquired:
		release()
		t.Fatal("second normalization entered while the only slot was held")
	case <-time.After(50 * time.Millisecond):
	}
	releaseFirst()
	select {
	case release := <-acquired:
		release()
	case <-time.After(time.Second):
		t.Fatal("second normalization did not acquire the released slot")
	}
}

func TestNormalizeVideoInputImagesPadsAspectRatioViolation(t *testing.T) {
	service := &Service{}
	source := pngImageDataURL(t, 1530, 500)
	body := map[string]any{"first_frame": source}

	normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
	if err != nil {
		t.Fatal(err)
	}
	if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(1530, 612) {
		t.Fatalf("unexpected converted size: %v", got)
	}
}

func TestNormalizeVideoInputImagesUpscalesBelowMinimumResolution(t *testing.T) {
	service := &Service{}
	body := map[string]any{"first_frame": pngImageDataURL(t, 200, 100)}

	normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
	if err != nil {
		t.Fatal(err)
	}
	if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(720, 360) {
		t.Fatalf("unexpected converted size: %v", got)
	}
}

func TestNormalizeVideoInputImagesKeepsCompliantImage(t *testing.T) {
	service := &Service{}
	source := pngImageDataURL(t, 1280, 720)
	body := map[string]any{"first_frame": source}

	normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
	if err != nil {
		t.Fatal(err)
	}
	if normalized["first_frame"] != source {
		t.Fatal("compliant image should remain unchanged")
	}
}

func TestNormalizeVideoInputImagesReturnsActionableDecodeError(t *testing.T) {
	service := &Service{}
	body := map[string]any{"first_frame": "data:image/png;base64,bm90LWltYWdl"}

	_, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
	if err == nil {
		t.Fatal("expected decode error")
	}
	if clients.ErrorCode(err) != "invalid_parameter" || clients.ErrorParam(err) != "first_frame" {
		t.Fatalf("unexpected error contract: code=%s param=%s err=%v", clients.ErrorCode(err), clients.ErrorParam(err), err)
	}
}