laimiaoai-agent-new/tests/gateway.test.ts

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<string, string> = {}) {
  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<Response> => {
      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<string, unknown>;
    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<Response> => {
      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<Response> => {
      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<Response> =>
        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<string, unknown>;
    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<Response> => 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<Response> => 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);
  });
});