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 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{"aspect_ratio": "4:5"}); got != "4:5" {
		t.Fatalf("expected explicit aspect ratio to win, 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)
	}
}