moteva/server/internal/infrastructure/piagent/image_generation_test.go

package piagent

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

func TestNewImageClientDefaultsToTenMinuteTimeout(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL: "http://example.com",
		APIKey:  "test-key",
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.client.Timeout != 10*time.Minute {
		t.Fatalf("expected 10 minute timeout, got %s", client.client.Timeout)
	}
}

func TestNewImageClientUsesConfiguredTimeout(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL:        "http://example.com",
		APIKey:         "test-key",
		TimeoutSeconds: 180,
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.client.Timeout != 3*time.Minute {
		t.Fatalf("expected configured 3 minute timeout, got %s", client.client.Timeout)
	}
}

func TestNewImageClientCapsConfiguredTimeoutAtTenMinutes(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL:        "http://example.com",
		APIKey:         "test-key",
		TimeoutSeconds: 900,
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.client.Timeout != 10*time.Minute {
		t.Fatalf("expected configured timeout to be capped at 10 minutes, got %s", client.client.Timeout)
	}
}

func TestNewImageClientDisablesKeepAliveReuse(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL: "http://example.com",
		APIKey:  "test-key",
	})
	transport, ok := client.client.Transport.(*http.Transport)
	if !ok {
		t.Fatalf("expected custom http transport, got %T", client.client.Transport)
	}
	if !transport.DisableKeepAlives {
		t.Fatal("expected image client to disable keep-alive reuse")
	}
}

func TestNewImageClientDefaultsInputImagesToFileTransport(t *testing.T) {
	client := NewImageClient(ImageClientOptions{
		BaseURL: "http://example.com",
		APIKey:  "test-key",
	})
	if client == nil {
		t.Fatal("expected image client")
	}
	if client.inputImageTransport != "file" {
		t.Fatalf("expected file input image transport, got %q", client.inputImageTransport)
	}
}

func TestImageClientSendsTransportDefaultSizeWhenUnset(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}); err != nil {
		t.Fatal(err)
	}
	if requestBody["size"] != "1024x1024" {
		t.Fatalf("expected transport default size, got %#v", requestBody["size"])
	}
	if _, ok := requestBody["n"]; ok {
		t.Fatalf("expected single image request to omit n for compatible transports, got %#v", requestBody["n"])
	}
}

func TestImageClientSendsIdempotencyKey(t *testing.T) {
	var idempotencyKey string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		idempotencyKey = r.Header.Get("Idempotency-Key")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", IdempotencyKey: "project-thread-image-0"}); err != nil {
		t.Fatal(err)
	}
	if idempotencyKey != "project-thread-image-0" {
		t.Fatalf("expected idempotency key header, got %q", idempotencyKey)
	}
}

func TestImageClientSendsExplicitSize(t *testing.T) {
	var requestBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", Size: "2048x1152"}); err != nil {
		t.Fatal(err)
	}
	if requestBody["size"] != "2048x1152" {
		t.Fatalf("expected explicit size, got %#v", requestBody["size"])
	}
}

func TestImageClientUploadsDataReferenceImagesAsMultipart(t *testing.T) {
	var requestPath string
	var requestContentType string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestPath = r.URL.Path
		requestContentType = r.Header.Get("Content-Type")
		if err := r.ParseMultipartForm(32 << 20); err != nil {
			t.Errorf("parse multipart form: %v", err)
			return
		}
		file, header, err := r.FormFile("image")
		if err != nil {
			t.Errorf("expected image file field: %v", err)
			return
		}
		defer file.Close()
		data, err := io.ReadAll(file)
		if err != nil {
			t.Errorf("read image file: %v", err)
			return
		}
		if header.Filename != "reference-1.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "reference bytes" {
			t.Errorf("unexpected uploaded image: filename=%q content_type=%q data=%q", header.Filename, header.Header.Get("Content-Type"), string(data))
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	_, err := client.Generate(context.Background(), "extract product from reference", ImageRequestOptions{
		Model:  "gpt-image-2",
		Images: []string{"data:image/png;base64,cmVmZXJlbmNlIGJ5dGVz"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if requestPath != "/v1/images/edits" {
		t.Fatalf("expected image edit endpoint for references, got %s", requestPath)
	}
	if !strings.HasPrefix(requestContentType, "multipart/form-data") {
		t.Fatalf("expected multipart edit request, got %s", requestContentType)
	}
}

func TestImageClientSendsReferenceImagesByURLWhenConfigured(t *testing.T) {
	var requestBody map[string]any
	var requestPath string
	var requestContentType string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestPath = r.URL.Path
		requestContentType = r.Header.Get("Content-Type")
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"})
	_, err := client.Generate(context.Background(), "extract product from reference", ImageRequestOptions{
		Model:  "gpt-image-2",
		Images: []string{" https://example.com/reference.png ", "", "https://example.com/reference.png", "https://example.com/second.webp"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if requestPath != "/v1/images/edits" {
		t.Fatalf("expected image edit endpoint for references, got %s", requestPath)
	}
	if requestContentType != "application/json" {
		t.Fatalf("expected url transport to use JSON, got %s", requestContentType)
	}
	images, ok := requestBody["images"].([]any)
	if !ok {
		t.Fatalf("expected images array in request, got %#v", requestBody["images"])
	}
	if len(images) != 2 {
		t.Fatalf("expected normalized reference images, got %#v", images)
	}
	first, ok := images[0].(map[string]any)
	if !ok || first["image_url"] != "https://example.com/reference.png" {
		t.Fatalf("expected first image_url reference, got %#v", images[0])
	}
	second, ok := images[1].(map[string]any)
	if !ok || second["image_url"] != "https://example.com/second.webp" {
		t.Fatalf("expected second image_url reference, got %#v", images[1])
	}
}

func TestImageClientUploadsInternalReferenceImagesAsMultipart(t *testing.T) {
	var sawEdit bool
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/reference.png":
			w.Header().Set("Content-Type", "image/png")
			_, _ = w.Write([]byte("internal image bytes"))
		case "/v1/images/edits":
			sawEdit = true
			if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") {
				t.Errorf("expected multipart edit request, got %s", r.Header.Get("Content-Type"))
			}
			if err := r.ParseMultipartForm(32 << 20); err != nil {
				t.Errorf("parse multipart form: %v", err)
				return
			}
			if r.FormValue("model") != "gpt-image-2" || r.FormValue("prompt") != "edit the internal image" || r.FormValue("size") != "1024x1024" {
				t.Errorf("unexpected multipart fields: model=%q prompt=%q size=%q", r.FormValue("model"), r.FormValue("prompt"), r.FormValue("size"))
			}
			file, header, err := r.FormFile("image")
			if err != nil {
				t.Errorf("expected image file field: %v", err)
				return
			}
			defer file.Close()
			data, err := io.ReadAll(file)
			if err != nil {
				t.Errorf("read image file: %v", err)
				return
			}
			if header.Filename != "reference.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "internal image bytes" {
				t.Errorf("unexpected uploaded image: filename=%q content_type=%q data=%q", header.Filename, header.Header.Get("Content-Type"), string(data))
			}
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/edited.png"}]}`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "edit the internal image", ImageRequestOptions{
		Model:  "gpt-image-2",
		Size:   "1024x1024",
		Images: []string{server.URL + "/reference.png"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/edited.png" {
		t.Fatalf("expected edited image, got %#v", image)
	}
	if !sawEdit {
		t.Fatal("expected image edit request")
	}
}

func TestImageClientParsesAlternateImageResponseShapes(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"output":[{"image_url":"https://example.com/generated.webp"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/generated.webp" {
		t.Fatalf("expected generated image url, got %#v", image)
	}
}

func TestImageClientRetriesTransientFailures(t *testing.T) {
	calls := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		if calls == 1 {
			http.Error(w, `{"error":"temporary overload"}`, http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}); err != nil {
		t.Fatal(err)
	}
	if calls != 2 {
		t.Fatalf("expected retry after transient failure, got %d calls", calls)
	}
}

func TestImageClientDoesNotRetryRequestTimeout(t *testing.T) {
	calls := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		time.Sleep(150 * time.Millisecond)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	client.client.Timeout = 20 * time.Millisecond
	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}); err == nil {
		t.Fatal("expected timeout error")
	}
	if calls != 1 {
		t.Fatalf("expected timeout to fail fast without retry, got %d calls", calls)
	}
}