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

package piagent

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"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 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 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)
	}
}