514 lines
17 KiB
JavaScript
514 lines
17 KiB
JavaScript
/**
|
|
* Unit tests for image generation handler
|
|
*
|
|
* Covers:
|
|
* - OpenAI-compatible format (openai, minimax, openrouter)
|
|
* - Gemini format (generateContent API)
|
|
* - Provider-specific formats (nanobanana, sdwebui)
|
|
* - Response normalization to OpenAI format
|
|
* - Error handling (missing prompt, invalid model)
|
|
*/
|
|
|
|
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
|
import { handleImageGenerationCore } from "../../open-sse/handlers/imageGenerationCore.js";
|
|
|
|
const originalFetch = global.fetch;
|
|
|
|
describe("handleImageGenerationCore", () => {
|
|
beforeEach(() => {
|
|
global.fetch = vi.fn();
|
|
});
|
|
|
|
afterEach(() => {
|
|
global.fetch = originalFetch;
|
|
vi.useRealTimers();
|
|
});
|
|
|
|
it("validates required prompt field", async () => {
|
|
const result = await handleImageGenerationCore({
|
|
body: { model: "openai/dall-e-3" },
|
|
modelInfo: { provider: "openai", model: "dall-e-3" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(false);
|
|
expect(result.status).toBe(400);
|
|
expect(result.error).toContain("Missing required field: prompt");
|
|
});
|
|
|
|
it("rejects unsupported provider", async () => {
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "test" },
|
|
modelInfo: { provider: "unknown-provider", model: "test" },
|
|
credentials: null,
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(false);
|
|
expect(result.status).toBe(400);
|
|
expect(result.error).toContain("does not support image generation");
|
|
});
|
|
|
|
it("generates image with OpenAI format", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
created: 1234567890,
|
|
data: [{ url: "https://example.com/image.png" }],
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A cute cat", n: 1, size: "1024x1024" },
|
|
modelInfo: { provider: "openai", model: "dall-e-3" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenCalledWith(
|
|
"https://api.openai.com/v1/images/generations",
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
headers: expect.objectContaining({
|
|
"Content-Type": "application/json",
|
|
Authorization: "Bearer test-key",
|
|
}),
|
|
body: expect.stringContaining('"prompt":"A cute cat"'),
|
|
})
|
|
);
|
|
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data).toHaveLength(1);
|
|
expect(responseBody.data[0].url).toBe("https://example.com/image.png");
|
|
});
|
|
|
|
it("generates image with Gemini format", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
candidates: [
|
|
{
|
|
content: {
|
|
parts: [
|
|
{ text: "Generated image" },
|
|
{ inlineData: { data: "base64imagedata" } },
|
|
],
|
|
},
|
|
},
|
|
],
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A sunset" },
|
|
modelInfo: { provider: "gemini", model: "gemini-image-preview" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenCalledWith(
|
|
expect.stringContaining("generativelanguage.googleapis.com"),
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
body: expect.stringContaining('"responseModalities":["TEXT","IMAGE"]'),
|
|
})
|
|
);
|
|
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data).toHaveLength(1);
|
|
expect(responseBody.data[0].b64_json).toBe("base64imagedata");
|
|
});
|
|
|
|
it("generates image with Minimax format", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
created: 1234567890,
|
|
data: [{ url: "https://example.com/minimax.png" }],
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A mountain", size: "1024x1024" },
|
|
modelInfo: { provider: "minimax", model: "minimax-image-01" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenCalledWith(
|
|
"https://api.minimaxi.com/v1/images/generations",
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
headers: expect.objectContaining({
|
|
Authorization: "Bearer test-key",
|
|
}),
|
|
})
|
|
);
|
|
});
|
|
|
|
it("generates image with NanoBanana format", async () => {
|
|
vi.useFakeTimers();
|
|
global.fetch
|
|
.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({ code: 200, data: { taskId: "task-123" } }),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
)
|
|
.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
data: {
|
|
successFlag: 1,
|
|
response: { resultImageUrl: "https://example.com/nanobanana.png" },
|
|
},
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const pending = handleImageGenerationCore({
|
|
body: { prompt: "A robot", n: 2, size: "1024x1792" },
|
|
modelInfo: { provider: "nanobanana", model: "nanobanana-flash" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
await vi.advanceTimersByTimeAsync(1500);
|
|
const result = await pending;
|
|
|
|
expect(result.success).toBe(true);
|
|
const fetchCall = global.fetch.mock.calls[0];
|
|
const requestBody = JSON.parse(fetchCall[1].body);
|
|
expect(requestBody.type).toBe("TEXTTOIAMGE");
|
|
expect(requestBody.numImages).toBe(2);
|
|
expect(requestBody.image_size).toBe("9:16");
|
|
expect(global.fetch).toHaveBeenNthCalledWith(
|
|
2,
|
|
"https://api.nanobananaapi.ai/api/v1/nanobanana/record-info?taskId=task-123",
|
|
expect.objectContaining({
|
|
headers: expect.objectContaining({
|
|
Authorization: "Bearer test-key",
|
|
}),
|
|
})
|
|
);
|
|
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data[0].url).toBe("https://example.com/nanobanana.png");
|
|
});
|
|
|
|
it("generates image with SD WebUI format", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({ images: ["base64sdwebui1", "base64sdwebui2"] }),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A forest", size: "768x768", n: 2 },
|
|
modelInfo: { provider: "sdwebui", model: "sdxl-base-1.0" },
|
|
credentials: null,
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
const fetchCall = global.fetch.mock.calls[0];
|
|
const requestBody = JSON.parse(fetchCall[1].body);
|
|
expect(requestBody.width).toBe(768);
|
|
expect(requestBody.height).toBe(768);
|
|
expect(requestBody.batch_size).toBe(2);
|
|
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data).toHaveLength(2);
|
|
});
|
|
|
|
it("handles OpenRouter with HTTP-Referer header", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
created: 1234567890,
|
|
data: [{ url: "https://example.com/or.png" }],
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A city" },
|
|
modelInfo: { provider: "openrouter", model: "openai/dall-e-3" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenCalledWith(
|
|
"https://openrouter.ai/api/v1/images/generations",
|
|
expect.objectContaining({
|
|
headers: expect.objectContaining({
|
|
"HTTP-Referer": "https://endpoint-proxy.local",
|
|
"X-Title": "Endpoint Proxy",
|
|
}),
|
|
})
|
|
);
|
|
});
|
|
|
|
it("handles HuggingFace binary response", async () => {
|
|
const imageBuffer = new Uint8Array([0x89, 0x50, 0x4e, 0x47]); // PNG header
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(imageBuffer, {
|
|
status: 200,
|
|
headers: { "Content-Type": "image/png" },
|
|
})
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A tree" },
|
|
modelInfo: { provider: "huggingface", model: "black-forest-labs/FLUX.1-schnell" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data[0].b64_json).toBeTruthy();
|
|
});
|
|
|
|
it("generates image with Codex gpt-5.5-image using current Codex version header", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
[
|
|
"event: response.output_item.done",
|
|
'data: {"item":{"type":"image_generation_call","result":"base64codeximage"}}',
|
|
"",
|
|
"",
|
|
].join("\n"),
|
|
{ status: 200, headers: { "Content-Type": "text/event-stream" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: {
|
|
prompt: "A green square",
|
|
size: "1024x1024",
|
|
output_format: "png",
|
|
},
|
|
modelInfo: { provider: "codex", model: "gpt-5.5-image" },
|
|
credentials: {
|
|
accessToken: "codex-token",
|
|
providerSpecificData: { chatgptAccountId: "account-123" },
|
|
},
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenCalledWith(
|
|
"https://chatgpt.com/backend-api/codex/responses",
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
headers: expect.objectContaining({
|
|
authorization: "Bearer codex-token",
|
|
"chatgpt-account-id": "account-123",
|
|
version: "0.129.0",
|
|
}),
|
|
})
|
|
);
|
|
|
|
const fetchCall = global.fetch.mock.calls[0];
|
|
const requestBody = JSON.parse(fetchCall[1].body);
|
|
expect(requestBody.model).toBe("gpt-5.5");
|
|
expect(requestBody.tools).toEqual([
|
|
{ type: "image_generation", output_format: "png", size: "1024x1024" },
|
|
]);
|
|
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data[0].b64_json).toBe("base64codeximage");
|
|
});
|
|
|
|
it("generates image with Cloudflare Workers AI JSON response", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
result: { image: "base64cloudflare" },
|
|
success: true,
|
|
errors: [],
|
|
messages: [],
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A lighthouse", size: "1024x1536" },
|
|
modelInfo: { provider: "cloudflare-ai", model: "@cf/leonardo/lucid-origin" },
|
|
credentials: {
|
|
apiKey: "cf-token",
|
|
providerSpecificData: { accountId: "cf-account" },
|
|
},
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenCalledWith(
|
|
"https://api.cloudflare.com/client/v4/accounts/cf-account/ai/run/@cf/leonardo/lucid-origin",
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
headers: expect.objectContaining({
|
|
"Content-Type": "application/json",
|
|
Authorization: "Bearer cf-token",
|
|
}),
|
|
})
|
|
);
|
|
|
|
const fetchCall = global.fetch.mock.calls[0];
|
|
const requestBody = JSON.parse(fetchCall[1].body);
|
|
expect(requestBody.prompt).toBe("A lighthouse");
|
|
expect(requestBody.width).toBe(1024);
|
|
expect(requestBody.height).toBe(1536);
|
|
|
|
const responseBody = await result.response.json();
|
|
expect(responseBody.data[0].b64_json).toBe("base64cloudflare");
|
|
});
|
|
|
|
it("uses multipart form data for Cloudflare FLUX.2 models", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
result: { image: "base64flux2" },
|
|
success: true,
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "A mountain lake", size: "1792x1024", steps: 4 },
|
|
modelInfo: { provider: "cloudflare-ai", model: "@cf/black-forest-labs/flux-2-klein-9b" },
|
|
credentials: {
|
|
apiKey: "cf-token",
|
|
providerSpecificData: { accountId: "cf-account" },
|
|
},
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
|
|
const fetchCall = global.fetch.mock.calls[0];
|
|
expect(fetchCall[1].headers).not.toHaveProperty("Content-Type");
|
|
expect(fetchCall[1].body).toBeInstanceOf(FormData);
|
|
expect(fetchCall[1].body.get("prompt")).toBe("A mountain lake");
|
|
expect(fetchCall[1].body.get("width")).toBe("1792");
|
|
expect(fetchCall[1].body.get("height")).toBe("1024");
|
|
expect(fetchCall[1].body.get("steps")).toBe("4");
|
|
});
|
|
|
|
it("resolves Cloudflare img2img and inpainting URL inputs before sending", async () => {
|
|
global.fetch
|
|
.mockResolvedValueOnce(new Response(new Uint8Array([1, 2, 3]), { status: 200, headers: { "Content-Type": "image/png" } }))
|
|
.mockResolvedValueOnce(new Response(new Uint8Array([4, 5, 6]), { status: 200, headers: { "Content-Type": "image/png" } }))
|
|
.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({ result: { image: "base64inpaint" }, success: true }),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: {
|
|
prompt: "Change to a lion",
|
|
image: "https://example.com/source.png",
|
|
mask_image: "https://example.com/mask.png",
|
|
size: "512x512",
|
|
},
|
|
modelInfo: { provider: "cloudflare-ai", model: "@cf/runwayml/stable-diffusion-v1-5-inpainting" },
|
|
credentials: {
|
|
apiKey: "cf-token",
|
|
providerSpecificData: { accountId: "cf-account" },
|
|
},
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(global.fetch).toHaveBeenNthCalledWith(1, "https://example.com/source.png");
|
|
expect(global.fetch).toHaveBeenNthCalledWith(2, "https://example.com/mask.png");
|
|
|
|
const providerCall = global.fetch.mock.calls[2];
|
|
expect(providerCall[0]).toBe("https://api.cloudflare.com/client/v4/accounts/cf-account/ai/run/@cf/runwayml/stable-diffusion-v1-5-inpainting");
|
|
const requestBody = JSON.parse(providerCall[1].body);
|
|
expect(requestBody.image).toEqual([1, 2, 3]);
|
|
expect(requestBody.image_b64).toBe(Buffer.from([1, 2, 3]).toString("base64"));
|
|
expect(requestBody.mask).toEqual([4, 5, 6]);
|
|
expect(requestBody.mask_image).toEqual([4, 5, 6]);
|
|
expect(requestBody.mask_b64).toBe(Buffer.from([4, 5, 6]).toString("base64"));
|
|
});
|
|
|
|
it("handles provider error responses", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({ error: { message: "Rate limit exceeded" } }),
|
|
{ status: 429, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "test" },
|
|
modelInfo: { provider: "openai", model: "dall-e-3" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(false);
|
|
expect(result.status).toBe(429);
|
|
expect(result.error).toContain("Rate limit exceeded");
|
|
});
|
|
|
|
it("handles network errors", async () => {
|
|
global.fetch.mockRejectedValueOnce(new Error("Network timeout"));
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "test" },
|
|
modelInfo: { provider: "openai", model: "dall-e-3" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
});
|
|
|
|
expect(result.success).toBe(false);
|
|
expect(result.status).toBe(502);
|
|
expect(result.error).toContain("Network timeout");
|
|
});
|
|
|
|
it("calls onRequestSuccess callback on success", async () => {
|
|
global.fetch.mockResolvedValueOnce(
|
|
new Response(
|
|
JSON.stringify({
|
|
created: 1234567890,
|
|
data: [{ url: "https://example.com/success.png" }],
|
|
}),
|
|
{ status: 200, headers: { "Content-Type": "application/json" } }
|
|
)
|
|
);
|
|
|
|
const onRequestSuccess = vi.fn();
|
|
|
|
const result = await handleImageGenerationCore({
|
|
body: { prompt: "test" },
|
|
modelInfo: { provider: "openai", model: "dall-e-3" },
|
|
credentials: { apiKey: "test-key" },
|
|
log: null,
|
|
onRequestSuccess,
|
|
});
|
|
|
|
expect(result.success).toBe(true);
|
|
expect(onRequestSuccess).toHaveBeenCalledTimes(1);
|
|
});
|
|
});
|