import { describe, expect, it, vi, afterEach } from "vitest"; import { GatewayClient, GatewayError, extractImageUrls, normalizeSize } from "../src/gateway/client.js"; import { silentLogger, tempRoot, testConfig } from "./helpers.js"; afterEach(() => { vi.unstubAllGlobals(); }); function client(overrides: Record = {}) { const config = testConfig(tempRoot(), overrides); return new GatewayClient(config.gateway, silentLogger); } describe("GatewayClient", () => { it("routes known CN model families to CN and unknown models to global", () => { const gateway = client({ GATEWAY_MOCK: "false", GATEWAY_CN_URL: "https://cn.test/", GATEWAY_GLOBAL_URL: "https://global.test/" }); expect(gateway.baseUrlFor("seedream@5.0")).toBe("https://cn.test"); expect(gateway.baseUrlFor("kling-v2")).toBe("https://cn.test"); expect(gateway.baseUrlFor("gpt-image-2")).toBe("https://global.test"); expect(gateway.baseUrlFor("unknown-model")).toBe("https://global.test"); }); it("honours GATEWAY_FORCE_URL", () => { const gateway = client({ GATEWAY_FORCE_URL: "https://forced.test/", GATEWAY_MOCK: "false" }); expect(gateway.baseUrlFor("seedream@5.0")).toBe("https://forced.test"); expect(gateway.baseUrlFor("gpt-image-2")).toBe("https://forced.test"); }); it("normalizes gpt-image sizes to multiples of 16", () => { expect(normalizeSize("gpt-image-2", "1024x1536")).toBe("1024x1536"); expect(normalizeSize("gpt-image-2", "1000x1501")).toBe("1008x1504"); expect(normalizeSize("seedream@5.0", "1000x1501")).toBe("1000x1501"); }); it("extracts URLs from common gateway response shapes", () => { expect( extractImageUrls({ data: [{ url: "https://a/1.png" }, { image_url: "https://a/2.png" }], output: { images: ["https://a/3.png"] }, drawing: "https://a/4.png", media_list: [{ oss_url: "https://a/5.png" }], }), ).toEqual(["https://a/1.png", "https://a/2.png", "https://a/3.png", "https://a/4.png", "https://a/5.png"]); }); it("generates mock URLs for each requested image", async () => { const gateway = client(); const result = await gateway.generateImage({ uid: "1", team_id: "1", model: "seedream@5.0", prompt: "a cat", size: "1024x1024", n: 2, }); expect(result.taskId).toMatch(/^mock-/); expect(result.urls).toHaveLength(2); expect(result.urls[0]).toContain("picsum.photos"); }); it("submits a task and polls until images are available", async () => { const gateway = client({ GATEWAY_MOCK: "false", GATEWAY_CN_URL: "https://cn.test", GATEWAY_POLL_INTERVAL_MS: "10", }); let call = 0; const fetchMock = vi.fn(async (_url: string | URL | Request, _init?: RequestInit): Promise => { call += 1; if (call === 1) { return new Response(JSON.stringify({ ret: 0, task_id: "task-1" }), { status: 200 }); } return new Response( JSON.stringify({ status: "running", data: call === 2 ? [] : [{ url: "https://img/1.png" }] }), { status: 200 }, ); }); vi.stubGlobal("fetch", fetchMock); const result = await gateway.generateImage({ uid: "1", team_id: "2", model: "seedream@5.0", prompt: "a cat", size: "1024x1536", n: 1, watermark: true, }); expect(result.taskId).toBe("task-1"); expect(result.urls).toEqual(["https://img/1.png"]); expect(fetchMock).toHaveBeenCalledTimes(3); expect(String(fetchMock.mock.calls[0]?.[0])).toBe("https://cn.test/generate/v1/images"); const body = JSON.parse(String(fetchMock.mock.calls[0]?.[1]?.body)) as Record; expect(body).toMatchObject({ uid: 1, team_id: 2, model: "seedream@5.0", n: 1, watermark: true }); expect(String(fetchMock.mock.calls[1]?.[0])).toBe("https://cn.test/generate/v1/images/task-1"); }); it("retries transient poll network failures", async () => { const gateway = client({ GATEWAY_MOCK: "false", GATEWAY_CN_URL: "https://cn.test", GATEWAY_POLL_INTERVAL_MS: "10", }); let call = 0; const fetchMock = vi.fn(async (): Promise => { call += 1; if (call === 1) { return new Response(JSON.stringify({ ret: 0, task_id: "task-retry" }), { status: 200 }); } if (call === 2) { throw new TypeError("fetch failed"); } return new Response(JSON.stringify({ status: "success", data: [{ url: "https://img/retry.png" }] }), { status: 200, }); }); vi.stubGlobal("fetch", fetchMock); const result = await gateway.generateImage({ uid: "1", team_id: "1", model: "seedream@5.0", prompt: "a cat", size: "1024x1024", n: 1, }); expect(result.urls).toEqual(["https://img/retry.png"]); expect(call).toBe(3); }); it("waits briefly when the gateway reports success before urls are attached", async () => { const gateway = client({ GATEWAY_MOCK: "false", GATEWAY_CN_URL: "https://cn.test", GATEWAY_POLL_INTERVAL_MS: "10", }); let call = 0; const fetchMock = vi.fn(async (): Promise => { call += 1; if (call === 1) { return new Response(JSON.stringify({ ret: 0, task_id: "task-finalize" }), { status: 200 }); } if (call === 2) { return new Response(JSON.stringify({ status: "success", data: [] }), { status: 200 }); } return new Response(JSON.stringify({ status: "success", data: [{ url: "https://img/final.png" }] }), { status: 200, }); }); vi.stubGlobal("fetch", fetchMock); const result = await gateway.generateImage({ uid: "1", team_id: "1", model: "seedream@5.0", prompt: "a cat", size: "1024x1024", n: 1, }); expect(result.urls).toEqual(["https://img/final.png"]); expect(call).toBe(3); }); it("returns inline urls without polling", async () => { const gateway = client({ GATEWAY_MOCK: "false" }); const fetchMock = vi.fn( async (_url: string | URL | Request, _init?: RequestInit): Promise => new Response(JSON.stringify({ ret: 0, task_id: "t", images: ["https://img/1.png"] }), { status: 200 }), ); vi.stubGlobal("fetch", fetchMock); const result = await gateway.generateImage({ uid: "1", team_id: "1", model: "gpt-image-2", prompt: "a cat", size: "1000x1501", n: 2, }); expect(result.urls).toEqual(["https://img/1.png"]); expect(fetchMock).toHaveBeenCalledTimes(1); const body = JSON.parse(String(fetchMock.mock.calls[0]?.[1]?.body)) as Record; expect(body.size).toBe("1008x1504"); }); it("fails fast on a non-zero gateway ret", async () => { const gateway = client({ GATEWAY_MOCK: "false" }); vi.stubGlobal( "fetch", vi.fn(async (): Promise => new Response(JSON.stringify({ ret: 1, msg: "no quota" }), { status: 200 })), ); await expect( gateway.generateImage({ uid: "1", team_id: "1", model: "seedream@5.0", prompt: "a cat", size: "1024x1024", n: 1, }), ).rejects.toThrow(/no quota/); }); it("maps HTTP failures to GatewayError", async () => { const gateway = client({ GATEWAY_MOCK: "false" }); vi.stubGlobal("fetch", vi.fn(async (): Promise => new Response("down", { status: 503 }))); await expect( gateway.generateImage({ uid: "1", team_id: "1", model: "seedream@5.0", prompt: "a cat", size: "1024x1024", n: 1, }), ).rejects.toBeInstanceOf(GatewayError); }); });