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":   "1024x1024",
	})
	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 normalized 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":   "1024x1536",
	}, 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 normalized 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{"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)
	}
}

func TestResolveSeedreamImageSizeUsesSupportedPresetOrFallback(t *testing.T) {
	if got := resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{"resolution": "4K"}); got != "4K" {
		t.Fatalf("expected 4K for Seedream 4.5, got %q", got)
	}
	if got := resolveSeedreamImageSize("doubao-seedream-4-5-251128", map[string]any{"size": "1024x1536"}); got != "2K" {
		t.Fatalf("expected 2K fallback for pixel size, got %q", got)
	}
	if got := resolveSeedreamImageSize("doubao-seedream-5-0-260128", map[string]any{"resolution": "3K"}); got != "3K" {
		t.Fatalf("expected 3K for Seedream 5.0, got %q", got)
	}
}