1import assert from "node:assert/strict";
2import test, { type TestContext } from "node:test";
3 
4import type { CliArgs } from "../types.ts";
5import {
6  buildRequestBody,
7  buildZaiUrl,
8  extractImageFromResponse,
9  getDefaultModel,
10  getModelFamily,
11  parseAspectRatio,
12  parseSize,
13  resolveSizeForModel,
14  validateArgs,
15} from "./zai.ts";
16 
17function makeArgs(overrides: Partial<CliArgs> = {}): CliArgs {
18  return {
19    prompt: null,
20    promptFiles: [],
21    imagePath: null,
22    provider: null,
23    model: null,
24    aspectRatio: null,
25    size: null,
26    quality: null,
27    imageSize: null,
28    imageApiDialect: null,
29    referenceImages: [],
30    n: 1,
31    batchFile: null,
32    jobs: null,
33    json: false,
34    help: false,
35    ...overrides,
36  };
37}
38 
39function useEnv(
40  t: TestContext,
41  values: Record<string, string | null>,
42): void {
43  const previous = new Map<string, string | undefined>();
44  for (const [key, value] of Object.entries(values)) {
45    previous.set(key, process.env[key]);
46    if (value == null) {
47      delete process.env[key];
48    } else {
49      process.env[key] = value;
50    }
51  }
52 
53  t.after(() => {
54    for (const [key, value] of previous.entries()) {
55      if (value == null) {
56        delete process.env[key];
57      } else {
58        process.env[key] = value;
59      }
60    }
61  });
62}
63 
64test("Z.AI default model prefers env override and otherwise uses glm-image", (t) => {
65  useEnv(t, {
66    ZAI_IMAGE_MODEL: null,
67    BIGMODEL_IMAGE_MODEL: null,
68  });
69  assert.equal(getDefaultModel(), "glm-image");
70 
71  process.env.BIGMODEL_IMAGE_MODEL = "cogview-4-250304";
72  assert.equal(getDefaultModel(), "cogview-4-250304");
73});
74 
75test("Z.AI URL builder normalizes host, v4 base, and full endpoint inputs", (t) => {
76  useEnv(t, { ZAI_BASE_URL: "https://api.z.ai" });
77  assert.equal(buildZaiUrl(), "https://api.z.ai/api/paas/v4/images/generations");
78 
79  process.env.ZAI_BASE_URL = "https://proxy.example.com/api/paas/v4/";
80  assert.equal(buildZaiUrl(), "https://proxy.example.com/api/paas/v4/images/generations");
81 
82  process.env.ZAI_BASE_URL = "https://proxy.example.com/custom/images/generations";
83  assert.equal(buildZaiUrl(), "https://proxy.example.com/custom/images/generations");
84});
85 
86test("Z.AI model family and parsing helpers recognize documented formats", () => {
87  assert.equal(getModelFamily("glm-image"), "glm");
88  assert.equal(getModelFamily("cogview-4-250304"), "legacy");
89  assert.deepEqual(parseAspectRatio("16:9"), { width: 16, height: 9 });
90  assert.equal(parseAspectRatio("wide"), null);
91  assert.deepEqual(parseSize("1280x1280"), { width: 1280, height: 1280 });
92  assert.deepEqual(parseSize("1472*1088"), { width: 1472, height: 1088 });
93  assert.equal(parseSize("big"), null);
94});
95 
96test("Z.AI size resolution follows documented recommended ratios and validates custom sizes", () => {
97  assert.equal(
98    resolveSizeForModel("glm-image", makeArgs({ aspectRatio: "16:9", quality: "2k" })),
99    "1728x960",
100  );
101  assert.equal(
102    resolveSizeForModel("cogview-4-250304", makeArgs({ aspectRatio: "4:3", quality: "normal" })),
103    "1152x864",
104  );
105  assert.equal(
106    resolveSizeForModel("glm-image", makeArgs({ size: "1568x1056", quality: "2k" })),
107    "1568x1056",
108  );
109 
110  const uncommon = resolveSizeForModel(
111    "glm-image",
112    makeArgs({ aspectRatio: "5:2", quality: "normal" }),
113  );
114  const parsed = parseSize(uncommon);
115  assert.ok(parsed);
116  assert.ok(parsed.width % 32 === 0);
117  assert.ok(parsed.height % 32 === 0);
118  assert.ok(parsed.width * parsed.height <= 2 ** 22);
119 
120  assert.throws(
121    () => resolveSizeForModel("glm-image", makeArgs({ size: "1000x1000", quality: "2k" })),
122    /between 1024 and 2048/,
123  );
124  assert.throws(
125    () => resolveSizeForModel("glm-image", makeArgs({ size: "1280x1260", quality: "2k" })),
126    /divisible by 32/,
127  );
128  assert.throws(
129    () => resolveSizeForModel("cogview-4-250304", makeArgs({ size: "2048x2048", quality: "2k" })),
130    /must not exceed 2\^21 total pixels/,
131  );
132});
133 
134test("Z.AI validation rejects unsupported refs and multi-image requests", () => {
135  assert.throws(
136    () => validateArgs("glm-image", makeArgs({ referenceImages: ["ref.png"] })),
137    /text-to-image only/,
138  );
139  assert.throws(
140    () => validateArgs("glm-image", makeArgs({ n: 2 })),
141    /single image per request/,
142  );
143});
144 
145test("Z.AI request body maps skill quality and resolved size into provider fields", () => {
146  const body = buildRequestBody(
147    "A cinematic science poster",
148    "glm-image",
149    makeArgs({ aspectRatio: "4:3", quality: "normal" }),
150  );
151 
152  assert.deepEqual(body, {
153    model: "glm-image",
154    prompt: "A cinematic science poster",
155    quality: "standard",
156    size: "1472x1088",
157  });
158});
159 
160test("Z.AI response extraction downloads the returned image URL", async (t) => {
161  const originalFetch = globalThis.fetch;
162  t.after(() => {
163    globalThis.fetch = originalFetch;
164  });
165 
166  globalThis.fetch = async () =>
167    new Response(Uint8Array.from([1, 2, 3]), {
168      status: 200,
169      headers: { "Content-Type": "image/png" },
170    });
171 
172  const image = await extractImageFromResponse({
173    data: [{ url: "https://cdn.example.com/glm-image.png" }],
174  });
175  assert.deepEqual([...image], [1, 2, 3]);
176 
177  await assert.rejects(
178    () => extractImageFromResponse({ data: [{}] }),
179    /No image URL/,
180  );
181});
182