1import assert from "node:assert/strict";
2import fs from "node:fs/promises";
3import os from "node:os";
4import path from "node:path";
5import test, { type TestContext } from "node:test";
6 
7import type { CliArgs } from "../types.ts";
8import {
9  buildRequestBody,
10  extractImageFromResponse,
11  parseAspectRatio,
12  resolveReferenceImages,
13  resolveSize,
14  snapDim,
15  validateArgs,
16} from "./agnes.ts";
17 
18function useEnv(
19  t: TestContext,
20  values: Record<string, string | null>,
21): void {
22  const previous = new Map<string, string | undefined>();
23  for (const [key, value] of Object.entries(values)) {
24    previous.set(key, process.env[key]);
25    if (value == null) {
26      delete process.env[key];
27    } else {
28      process.env[key] = value;
29    }
30  }
31 
32  t.after(() => {
33    for (const [key, value] of previous.entries()) {
34      if (value == null) {
35        delete process.env[key];
36      } else {
37        process.env[key] = value;
38      }
39    }
40  });
41}
42 
43function makeArgs(overrides: Partial<CliArgs> = {}): CliArgs {
44  return {
45    prompt: null,
46    promptFiles: [],
47    imagePath: null,
48    provider: null,
49    model: null,
50    aspectRatio: null,
51    size: null,
52    quality: null,
53    imageSize: null,
54    imageApiDialect: null,
55    referenceImages: [],
56    n: 1,
57    batchFile: null,
58    jobs: null,
59    json: false,
60    help: false,
61    responseFormat: null,
62    ...overrides,
63  };
64}
65 
66test("snapDim rounds to the nearest multiple of 32", () => {
67  assert.equal(snapDim(767), 768);
68  assert.equal(snapDim(1023), 1024);
69  assert.equal(snapDim(1024), 1024);
70  assert.equal(snapDim(32), 32);
71  assert.equal(snapDim(0), 32);
72  assert.equal(snapDim(16), 32);
73  assert.equal(snapDim(48), 64);
74});
75 
76test("parseAspectRatio parses valid ratios and rejects invalid inputs", () => {
77  assert.deepEqual(parseAspectRatio("3:4"), { width: 3, height: 4 });
78  assert.deepEqual(parseAspectRatio("16:9"), { width: 16, height: 9 });
79  assert.deepEqual(parseAspectRatio("1:1"), { width: 1, height: 1 });
80  assert.deepEqual(parseAspectRatio("1.5:1"), { width: 1.5, height: 1 });
81 
82  assert.equal(parseAspectRatio(""), null);
83  assert.equal(parseAspectRatio("invalid"), null);
84  assert.equal(parseAspectRatio("3x4"), null);
85  assert.equal(parseAspectRatio("0:1"), null);
86  assert.equal(parseAspectRatio("1:0"), null);
87});
88 
89test("resolveSize returns explicit --size directly", () => {
90  assert.equal(resolveSize({ size: "1024x1024" }), "1024x1024");
91  assert.equal(resolveSize({ size: "768x1024", aspectRatio: "16:9" }), "768x1024");
92});
93 
94test("resolveSize returns default 1024x1024 when no size or ratio given", () => {
95  assert.equal(resolveSize({}), "1024x1024");
96  assert.equal(resolveSize({ size: null, aspectRatio: null }), "1024x1024");
97});
98 
99test("resolveSize computes 32-aligned size within 2048 max edge", () => {
100  assert.equal(resolveSize({ aspectRatio: "1:1" }), "1024x1024");
101  assert.equal(resolveSize({ aspectRatio: "16:9" }), "2048x1152");
102  assert.equal(resolveSize({ aspectRatio: "4:3" }), "2048x1536");
103  assert.equal(resolveSize({ aspectRatio: "3:4" }), "1536x2048");
104  assert.equal(resolveSize({ aspectRatio: "9:16" }), "1152x2048");
105});
106 
107test("resolveSize aligns to 32 and respects max edge", () => {
108  assert.equal(resolveSize({ aspectRatio: "3:1" }), "2048x672");
109  assert.equal(resolveSize({ aspectRatio: "1:3" }), "672x2048");
110});
111 
112test("validateArgs rejects --n > 1", () => {
113  assert.throws(
114    () => validateArgs("agnes-image-2.1-flash", makeArgs({ n: 2 })),
115    /returns a single image per request/,
116  );
117  assert.doesNotThrow(() =>
118    validateArgs("agnes-image-2.1-flash", makeArgs({ n: 1 })),
119  );
120});
121 
122test("buildRequestBody maps prompt, model, size, and reference images", () => {
123  const body = buildRequestBody("a cat", "agnes-image-2.1-flash", {
124    size: "1024x1024",
125    aspectRatio: null,
126    referenceImages: [],
127  });
128  assert.equal(body.model, "agnes-image-2.1-flash");
129  assert.equal(body.prompt, "a cat");
130  assert.equal(body.size, "1024x1024");
131  assert.deepEqual(body.extra_body, { response_format: "url" });
132 
133  const bodyWithRef = buildRequestBody("a cat", "agnes-image-2.1-flash", {
134    size: null,
135    aspectRatio: "3:4",
136    referenceImages: ["https://example.com/ref.jpg"],
137  });
138  assert.equal(bodyWithRef.size, "1536x2048");
139  assert.deepEqual(bodyWithRef.image, ["https://example.com/ref.jpg"]);
140});
141 
142test("extractImageFromResponse decodes b64_json payloads", async () => {
143  const fromBase64 = await extractImageFromResponse({
144    data: [{ b64_json: Buffer.from("hello").toString("base64") }],
145  });
146  assert.equal(Buffer.from(fromBase64).toString("utf8"), "hello");
147});
148 
149test("extractImageFromResponse downloads URL payloads", async (t) => {
150  const originalFetch = globalThis.fetch;
151  t.after(() => {
152    globalThis.fetch = originalFetch;
153  });
154 
155  globalThis.fetch = async () =>
156    new Response(Uint8Array.from([1, 2, 3]), {
157      status: 200,
158      headers: { "Content-Type": "image/png" },
159    });
160 
161  const fromUrl = await extractImageFromResponse({
162    data: [{ url: "https://example.com/output.png" }],
163  });
164  assert.deepEqual([...fromUrl], [1, 2, 3]);
165});
166 
167test("extractImageFromResponse throws on empty data", async () => {
168  await assert.rejects(
169    () => extractImageFromResponse({ data: [] }),
170    /No image/,
171  );
172  await assert.rejects(
173    () => extractImageFromResponse({ data: [{}] }),
174    /No image/,
175  );
176});
177 
178test("resolveReferenceImages converts local files to data URIs and passes URLs through", async (t) => {
179  const dir = await fs.mkdtemp(path.join(os.tmpdir(), "agnes-ref-"));
180  t.after(() => fs.rm(dir, { recursive: true, force: true }));
181 
182  const localPath = path.join(dir, "ref.png");
183  const localBytes = Buffer.from([0x89, 0x50, 0x4e, 0x47]);
184  await fs.writeFile(localPath, localBytes);
185 
186  const jpegPath = path.join(dir, "photo.jpeg");
187  await fs.writeFile(jpegPath, Buffer.from([0xff, 0xd8]));
188 
189  const results = await resolveReferenceImages([
190    localPath,
191    "https://example.com/remote.jpg",
192    jpegPath,
193  ]);
194 
195  assert.equal(results.length, 3);
196  assert.match(results[0]!, /^data:image\/png;base64,/);
197  assert.match(results[1]!, /^https:\/\/example.com\/remote.jpg$/);
198  assert.match(results[2]!, /^data:image\/jpeg;base64,/);
199});
200 
201test("resolveReferenceImages detects gif and webp mime types", async (t) => {
202  const dir = await fs.mkdtemp(path.join(os.tmpdir(), "agnes-mime-"));
203  t.after(() => fs.rm(dir, { recursive: true, force: true }));
204 
205  const webpPath = path.join(dir, "ref.webp");
206  const gifPath = path.join(dir, "ref.gif");
207  await fs.writeFile(webpPath, Buffer.from([0x00]));
208  await fs.writeFile(gifPath, Buffer.from([0x00]));
209 
210  const results = await resolveReferenceImages([webpPath, gifPath]);
211  assert.match(results[0]!, /^data:image\/webp;base64,/);
212  assert.match(results[1]!, /^data:image\/gif;base64,/);
213});
214