popiart-server/internal/server/newapi_image_test.go

package server

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestGenerateGeminiImageRefsUsesGenerateContentEndpoint(t *testing.T) {
	const png1x1 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z7xkAAAAASUVORK5CYII="

	var gotPath string
	var gotAuth string
	var gotBody map[string]any

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotAuth = r.Header.Get("Authorization")
		if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
			t.Fatalf("decode request body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"candidates": [{
				"content": {
					"parts": [{
						"inlineData": {
							"mimeType": "image/png",
							"data": "` + png1x1 + `"
						}
					}]
				}
			}],
			"usageMetadata": {"promptTokenCount": 12}
		}`))
	}))
	defer srv.Close()

	client := &newAPIClient{
		baseURL:    strings.TrimRight(srv.URL, "/"),
		httpClient: srv.Client(),
	}
	refBytes, err := base64.StdEncoding.DecodeString(png1x1)
	if err != nil {
		t.Fatalf("decode fixture: %v", err)
	}

	refs, usage, err := client.generateEditedImageRefs(context.Background(), "sk-test", "gemini-3-pro-image-preview", map[string]any{
		"prompt": "edit this image",
		"size":   "1024x1536",
	}, []imageEditReference{{
		Filename:    "source.png",
		ContentType: "image/png",
		Content:     refBytes,
	}})
	if err != nil {
		t.Fatalf("generateEditedImageRefs: %v", err)
	}
	if gotPath != "/v1beta/models/gemini-3-pro-image-preview:generateContent" {
		t.Fatalf("unexpected path: %s", gotPath)
	}
	if gotAuth != "Bearer sk-test" {
		t.Fatalf("unexpected auth header: %q", gotAuth)
	}
	if len(refs) != 1 || refs[0].Kind != "data_url" {
		t.Fatalf("expected one data_url result, got %#v", refs)
	}
	if usage["promptTokenCount"] != float64(12) {
		t.Fatalf("unexpected usage metadata: %#v", usage)
	}

	contents, ok := gotBody["contents"].([]any)
	if !ok || len(contents) != 1 {
		t.Fatalf("unexpected contents: %#v", gotBody["contents"])
	}
	content, ok := contents[0].(map[string]any)
	if !ok {
		t.Fatalf("unexpected content entry: %#v", contents[0])
	}
	parts, ok := content["parts"].([]any)
	if !ok || len(parts) != 2 {
		t.Fatalf("expected prompt plus one image part, got %#v", content["parts"])
	}

	generationConfig, ok := gotBody["generationConfig"].(map[string]any)
	if !ok {
		t.Fatalf("unexpected generationConfig: %#v", gotBody["generationConfig"])
	}
	imageConfig, ok := generationConfig["imageConfig"].(map[string]any)
	if !ok {
		t.Fatalf("expected imageConfig to exist, got %#v", generationConfig["imageConfig"])
	}
	if imageConfig["aspectRatio"] != "2:3" {
		t.Fatalf("unexpected aspectRatio: %#v", imageConfig["aspectRatio"])
	}
	if imageConfig["imageSize"] != "2K" {
		t.Fatalf("unexpected imageSize: %#v", imageConfig["imageSize"])
	}
}

func TestGenerateSeedreamImageRefsUsesImagesGenerationsEndpoint(t *testing.T) {
	var gotPath string
	var gotAuth string
	var gotBody map[string]any

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotAuth = r.Header.Get("Authorization")
		if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
			t.Fatalf("decode request body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"created": 1757388756,
			"data": [{"url":"https://example.com/generated.png","size":"2720x1536"}],
			"usage": {"generated_images": 1}
		}`))
	}))
	defer srv.Close()

	client := &newAPIClient{
		baseURL:    strings.TrimRight(srv.URL, "/"),
		httpClient: srv.Client(),
	}

	refs, usage, err := client.generateImageRefs(context.Background(), "sk-test", "seedream-4-5-251128", map[string]any{
		"prompt": "draw a golden retriever in the park",
		"size":   "2K",
	})
	if err != nil {
		t.Fatalf("generateImageRefs: %v", err)
	}
	if gotPath != "/v1/images/generations" {
		t.Fatalf("unexpected path: %s", gotPath)
	}
	if gotAuth != "Bearer sk-test" {
		t.Fatalf("unexpected auth header: %q", gotAuth)
	}
	if gotBody["size"] != "2K" {
		t.Fatalf("expected preset 2K size, got %#v", gotBody["size"])
	}
	if gotBody["response_format"] != "url" {
		t.Fatalf("expected default response_format=url, got %#v", gotBody["response_format"])
	}
	if _, exists := gotBody["image"]; exists {
		t.Fatalf("did not expect image field for text2image payload: %#v", gotBody["image"])
	}
	if len(refs) != 1 || refs[0].Kind != "url" || refs[0].URL != "https://example.com/generated.png" {
		t.Fatalf("unexpected refs: %#v", refs)
	}
	if usage["generated_images"] != float64(1) {
		t.Fatalf("unexpected usage: %#v", usage)
	}
}

func TestGenerateEditedImageRefsUsesSeedreamImagesGenerationsEndpoint(t *testing.T) {
	var gotPath string
	var gotAuth string
	var gotBody map[string]any

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotAuth = r.Header.Get("Authorization")
		if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
			t.Fatalf("decode request body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"created": 1757388756,
			"data": [{"url":"https://example.com/edited.png","size":"2720x1536"}],
			"usage": {"generated_images": 1}
		}`))
	}))
	defer srv.Close()

	client := &newAPIClient{
		baseURL:    strings.TrimRight(srv.URL, "/"),
		httpClient: srv.Client(),
	}

	refs, usage, err := client.generateEditedImageRefs(context.Background(), "sk-test", "seedream-4-5-251128", map[string]any{
		"prompt": "edit this image into a dusk scene",
		"size":   "2K",
	}, []imageEditReference{{
		Filename:    "source.jpg",
		ContentType: "image/jpeg",
		Content:     []byte("binary-image"),
		URL:         "https://example.com/reference.jpg",
	}})
	if err != nil {
		t.Fatalf("generateEditedImageRefs: %v", err)
	}
	if gotPath != "/v1/images/generations" {
		t.Fatalf("unexpected path: %s", gotPath)
	}
	if gotAuth != "Bearer sk-test" {
		t.Fatalf("unexpected auth header: %q", gotAuth)
	}
	if gotBody["size"] != "2K" {
		t.Fatalf("expected preset 2K size, got %#v", gotBody["size"])
	}
	if gotBody["image"] != "https://example.com/reference.jpg" {
		t.Fatalf("expected image URL payload, got %#v", gotBody["image"])
	}
	if len(refs) != 1 || refs[0].Kind != "url" || refs[0].URL != "https://example.com/edited.png" {
		t.Fatalf("unexpected refs: %#v", refs)
	}
	if usage["generated_images"] != float64(1) {
		t.Fatalf("unexpected usage: %#v", usage)
	}
}

func TestResolveGeminiAspectRatioFromSize(t *testing.T) {
	if got := resolveGeminiAspectRatio(map[string]any{"size": "1792x1024"}); got != "16:9" {
		t.Fatalf("expected 16:9, got %q", got)
	}
	if got := resolveGeminiAspectRatio(map[string]any{"size": "2520x1080"}); got != "21:9" {
		t.Fatalf("expected 21:9 from exact size, got %q", got)
	}
	if got := resolveGeminiAspectRatio(map[string]any{"aspect_ratio": "4:5"}); got != "4:5" {
		t.Fatalf("expected explicit aspect ratio to win, got %q", got)
	}
	if got := resolveGeminiAspectRatio(map[string]any{"aspect_ratio": "4x5"}); got != "4:5" {
		t.Fatalf("expected normalized 4:5, got %q", got)
	}
}

func TestResolveGeminiImageSizeFromResolutionOrSize(t *testing.T) {
	if got := resolveGeminiImageSize(map[string]any{"resolution": "4K"}); got != "4K" {
		t.Fatalf("expected explicit 4K, got %q", got)
	}
	if got := resolveGeminiImageSize(map[string]any{"size": "1024x1024"}); got != "1K" {
		t.Fatalf("expected 1K from square preset, got %q", got)
	}
	if got := resolveGeminiImageSize(map[string]any{"size": "2520x1080"}); got != "2K" {
		t.Fatalf("expected 2K bucket for 2520x1080, got %q", got)
	}
}

func TestResolveSeedreamImageSizeUsesSupportedPresetOrFallback(t *testing.T) {
	got, err := resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{"resolution": "4K"})
	if err != nil {
		t.Fatalf("expected no error for 4K preset, got %v", err)
	}
	if got != "4K" {
		t.Fatalf("expected 4K for Seedream 4.5, got %q", got)
	}
	if _, err := resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{"size": "1024x1536"}); err == nil {
		t.Fatal("expected exact pixel size to be rejected for Seedream")
	}
	got, err = resolveSeedreamImageSize("doubao-seedream-5-0-260128", map[string]any{"resolution": "3K"})
	if err != nil {
		t.Fatalf("expected no error for 3K preset, got %v", err)
	}
	if got != "3K" {
		t.Fatalf("expected 3K for Seedream 5.0, got %q", got)
	}
	if _, err := resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{"size": "2048x1152"}); err == nil {
		t.Fatal("expected exact size passthrough to be rejected for Seedream")
	}
	if _, err := resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{"aspect_ratio": "4:5"}); err == nil {
		t.Fatal("expected aspect_ratio to be rejected for Seedream")
	}
	got, err = resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{})
	if err != nil {
		t.Fatalf("expected default 2K without explicit size, got %v", err)
	}
	if got != "2K" {
		t.Fatalf("expected default 2K, got %q", got)
	}
}

func TestGenerateMiniMaxImageRefsUsesImagesGenerationsEndpoint(t *testing.T) {
	var gotPath string
	var gotBody map[string]any

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
			t.Fatalf("decode request body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"created": 1711234567,
			"data": [{"url":"https://example.com/minimax.png"}]
		}`))
	}))
	defer srv.Close()

	client := &newAPIClient{
		baseURL:    strings.TrimRight(srv.URL, "/"),
		httpClient: srv.Client(),
	}

	refs, _, err := client.generateImageRefs(context.Background(), "sk-test", "image-01", map[string]any{
		"prompt":           "draw a portrait",
		"aspect_ratio":     "4x5",
		"response_format":  "b64_json",
		"prompt_optimizer": true,
		"watermark":        false,
	})
	if err != nil {
		t.Fatalf("generateImageRefs: %v", err)
	}
	if gotPath != "/v1/images/generations" {
		t.Fatalf("unexpected path: %s", gotPath)
	}
	if gotBody["aspect_ratio"] != "4:5" {
		t.Fatalf("expected normalized aspect ratio 4:5, got %#v", gotBody["aspect_ratio"])
	}
	if gotBody["response_format"] != "b64_json" {
		t.Fatalf("expected b64_json response format, got %#v", gotBody["response_format"])
	}
	if gotBody["prompt_optimizer"] != true {
		t.Fatalf("expected prompt_optimizer=true, got %#v", gotBody["prompt_optimizer"])
	}
	if len(refs) != 1 || refs[0].URL != "https://example.com/minimax.png" {
		t.Fatalf("unexpected refs: %#v", refs)
	}
}

func TestGenerateMiniMaxEditedImageRefsSendsStandardImageField(t *testing.T) {
	var gotBody map[string]any

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
			t.Fatalf("decode request body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"created": 1711234567,
			"data": [{"url":"https://example.com/minimax-edit.png"}]
		}`))
	}))
	defer srv.Close()

	client := &newAPIClient{
		baseURL:    strings.TrimRight(srv.URL, "/"),
		httpClient: srv.Client(),
	}

	_, _, err := client.generateEditedImageRefs(context.Background(), "sk-test", "image-01", map[string]any{
		"prompt":          "turn it into watercolor",
		"size":            "832x1248",
		"response_format": "url",
	}, []imageEditReference{{
		URL: "https://example.com/reference.jpg",
	}})
	if err != nil {
		t.Fatalf("generateEditedImageRefs: %v", err)
	}
	if gotBody["image"] != "https://example.com/reference.jpg" {
		t.Fatalf("expected standard image field, got %#v", gotBody["image"])
	}
	if gotBody["size"] != "832x1248" {
		t.Fatalf("expected size passthrough for minimax mapping, got %#v", gotBody["size"])
	}
}

func TestGenerateGeminiImageRefsAnnotatesMultiImageRoles(t *testing.T) {
	const png1x1 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z7xkAAAAASUVORK5CYII="

	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
			t.Fatalf("decode request body: %v", err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"candidates": [{
				"content": {
					"parts": [{
						"inlineData": {
							"mimeType": "image/png",
							"data": "` + png1x1 + `"
						}
					}]
				}
			}]
		}`))
	}))
	defer srv.Close()

	refBytes, err := base64.StdEncoding.DecodeString(png1x1)
	if err != nil {
		t.Fatalf("decode fixture: %v", err)
	}
	client := &newAPIClient{
		baseURL:    strings.TrimRight(srv.URL, "/"),
		httpClient: srv.Client(),
	}

	_, _, err = client.generateGeminiImageRefs(context.Background(), "sk-test", "gemini-3-pro-image-preview", map[string]any{
		"prompt": "Replace the person in Image 1 with the character from Image 2 and apply the style from Image 3.",
	}, []imageEditReference{
		{Role: "source", Filename: "source.png", ContentType: "image/png", Content: refBytes},
		{Role: "identity", Filename: "identity.png", ContentType: "image/png", Content: refBytes},
		{Role: "style", Filename: "style.png", ContentType: "image/png", Content: refBytes},
	})
	if err != nil {
		t.Fatalf("generateGeminiImageRefs: %v", err)
	}

	contents, ok := gotBody["contents"].([]any)
	if !ok || len(contents) != 1 {
		t.Fatalf("unexpected contents: %#v", gotBody["contents"])
	}
	content, ok := contents[0].(map[string]any)
	if !ok {
		t.Fatalf("unexpected content entry: %#v", contents[0])
	}
	parts, ok := content["parts"].([]any)
	if !ok || len(parts) != 7 {
		t.Fatalf("expected prompt plus three labeled images, got %#v", content["parts"])
	}
	if got := parts[1].(map[string]any)["text"]; got != "Image 1 is the source scene. Keep the scene layout, camera framing, main action, and spatial relationships from this image." {
		t.Fatalf("unexpected source label: %#v", got)
	}
	if got := parts[3].(map[string]any)["text"]; got != "Image 2 is the identity reference. Keep the character face, hair, accessories, and recognizability from this image." {
		t.Fatalf("unexpected identity label: %#v", got)
	}
	if got := parts[5].(map[string]any)["text"]; got != "Image 3 is the style reference. Apply only the visual style, palette, texture, and illustration treatment from this image. Do not change the character identity because of this image." {
		t.Fatalf("unexpected style label: %#v", got)
	}
}

func TestDecodeGeminiImageGenerationResponseAcceptsStringErrorCode(t *testing.T) {
	_, _, err := decodeGeminiImageGenerationResponse(http.StatusBadRequest, []byte(`{"error":{"code":"model_not_found","message":"model unavailable","status":"INVALID_ARGUMENT"}}`), "gemini-3.1-flash-image-preview")
	if err == nil {
		t.Fatal("expected Gemini error response to fail")
	}
	if !strings.Contains(err.Error(), "model unavailable") {
		t.Fatalf("unexpected error: %v", err)
	}
}