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