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 requestBody["stream"] != true {
		t.Fatalf("expected image generations to use stream transport, got %#v", requestBody["stream"])
	}
	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 TestImageClientUploadsLocalhostReferenceImagesEvenWhenURLTransportConfigured(t *testing.T) {
	var sawMultipartEdit 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("local image bytes"))
		case "/v1/images/edits":
			sawMultipartEdit = true
			if !strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data") {
				t.Errorf("expected multipart edit request for local image, got %s", 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.png" || header.Header.Get("Content-Type") != "image/png" || string(data) != "local 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", InputImageTransport: "url"})
	image, err := client.Generate(context.Background(), "edit the local image", ImageRequestOptions{
		Model:  "gpt-image-2",
		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 !sawMultipartEdit {
		t.Fatal("expected multipart image edit request")
	}
}

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 TestImageClientParsesStreamingImageResponse(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", "text/event-stream")
		_, _ = w.Write([]byte("event: image_generation.completed\n"))
		_, _ = w.Write([]byte(`data: {"type":"image_generation.completed","result":"Z2VuZXJhdGVkIGJ5dGVz"}` + "\n\n"))
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
	}))
	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 requestBody["stream"] != true {
		t.Fatalf("expected stream request body, got %#v", requestBody)
	}
	if image.URL != "data:image/png;base64,Z2VuZXJhdGVkIGJ5dGVz" {
		t.Fatalf("expected streamed base64 image, got %#v", image)
	}
}

func TestImageClientFallsBackWhenStreamIsUnsupported(t *testing.T) {
	calls := 0
	streamValues := make([]any, 0, 2)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		var requestBody map[string]any
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		streamValues = append(streamValues, requestBody["stream"])
		if calls == 1 {
			http.Error(w, `{"error":{"message":"unknown parameter: stream"}}`, http.StatusBadRequest)
			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"})
	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/image.png" {
		t.Fatalf("expected fallback image, got %#v", image)
	}
	if calls != 2 {
		t.Fatalf("expected stream then json fallback, got %d calls", calls)
	}
	if len(streamValues) != 2 || streamValues[0] != true || streamValues[1] != nil {
		t.Fatalf("expected first request stream=true and fallback without stream, got %#v", streamValues)
	}
}

func TestImageClientFallsBackAfterIdempotentStreamFailures(t *testing.T) {
	calls := 0
	streamValues := make([]any, 0, 4)
	keys := make([]string, 0, 4)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		keys = append(keys, r.Header.Get("Idempotency-Key"))
		var requestBody map[string]any
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatal(err)
		}
		streamValues = append(streamValues, requestBody["stream"])
		if calls <= 3 {
			http.Error(w, `{"error":{"message":"temporary stream failure"}}`, http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`))
	}))
	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",
		IdempotencyKey: "same-generation",
	})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/recovered.png" {
		t.Fatalf("expected fallback image, got %#v", image)
	}
	if calls != 4 {
		t.Fatalf("expected three stream attempts then json fallback, got %d calls", calls)
	}
	if len(streamValues) != 4 || streamValues[0] != true || streamValues[1] != true || streamValues[2] != true || streamValues[3] != nil {
		t.Fatalf("expected stream retries then json fallback, got %#v", streamValues)
	}
	for _, key := range keys {
		if key != "same-generation" {
			t.Fatalf("expected idempotency key to be preserved, got %#v", keys)
		}
	}
}

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)
	}
}

func TestImageClientRetriesIdempotentResponseHeaderTimeout(t *testing.T) {
	calls := 0
	var keys []string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		keys = append(keys, r.Header.Get("Idempotency-Key"))
		if calls == 1 {
			time.Sleep(80 * time.Millisecond)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`))
	}))
	defer server.Close()

	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
	transport, ok := client.client.Transport.(*http.Transport)
	if !ok {
		t.Fatalf("expected http transport, got %T", client.client.Transport)
	}
	transport.ResponseHeaderTimeout = 20 * time.Millisecond
	client.client.Timeout = time.Second

	image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", IdempotencyKey: "same-generation"})
	if err != nil {
		t.Fatal(err)
	}
	if image.URL != "https://example.com/recovered.png" {
		t.Fatalf("expected recovered image, got %#v", image)
	}
	if calls != 2 {
		t.Fatalf("expected one retry after timeout, got %d calls", calls)
	}
	if len(keys) != 2 || keys[0] != "same-generation" || keys[1] != "same-generation" {
		t.Fatalf("expected retry to keep idempotency key, got %#v", keys)
	}
}