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  buildImageInput,
10  buildRequestBody,
11  generateImage,
12  getDefaultOutputExtension,
13  resolveSeedreamSize,
14  validateArgs,
15} from "./seedream.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 
64async function makeTempPng(t: TestContext, name: string): Promise<string> {
65  const dir = await fs.mkdtemp(path.join(os.tmpdir(), "seedream-test-"));
66  t.after(() => fs.rm(dir, { recursive: true, force: true }));
67 
68  const filePath = path.join(dir, name);
69  const png1x1 =
70    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+a7m0AAAAASUVORK5CYII=";
71  await fs.writeFile(filePath, Buffer.from(png1x1, "base64"));
72  return filePath;
73}
74 
75test("Seedream request body and default extensions follow official model capabilities", () => {
76  const five = buildRequestBody(
77    "A robot illustrator",
78    "doubao-seedream-5-0-260128",
79    makeArgs(),
80  );
81  assert.equal(five.size, "2K");
82  assert.equal(five.response_format, "url");
83  assert.equal(five.output_format, "png");
84  assert.equal(getDefaultOutputExtension("doubao-seedream-5-0-260128"), ".png");
85 
86  const fourFive = buildRequestBody(
87    "A robot illustrator",
88    "doubao-seedream-4-5-251128",
89    makeArgs(),
90  );
91  assert.equal(fourFive.size, "2K");
92  assert.equal(fourFive.response_format, "url");
93  assert.ok(!("output_format" in fourFive));
94  assert.equal(getDefaultOutputExtension("doubao-seedream-4-5-251128"), ".jpg");
95 
96  assert.throws(
97    () =>
98      buildRequestBody(
99        "Change the bubbles into hearts",
100        "doubao-seededit-3-0-i2i-250628",
101        makeArgs({ referenceImages: ["ref.png"] }),
102        "data:image/png;base64,AAAA",
103      ),
104    /no longer supported/,
105  );
106});
107 
108test("Seedream size selection validates model-specific presets", () => {
109  assert.equal(
110    resolveSeedreamSize("doubao-seedream-4-0-250828", makeArgs({ quality: "normal" })),
111    "1K",
112  );
113  assert.equal(
114    resolveSeedreamSize("doubao-seedream-3-0-t2i-250415", makeArgs({ quality: "2k" })),
115    "2048x2048",
116  );
117 
118  assert.throws(
119    () =>
120      resolveSeedreamSize("doubao-seedream-5-0-260128", makeArgs({ size: "4K" })),
121    /only supports 2K, 3K/,
122  );
123  assert.throws(
124    () =>
125      resolveSeedreamSize("doubao-seedream-3-0-t2i-250415", makeArgs({ imageSize: "2K" })),
126    /only supports explicit WxH sizes/,
127  );
128  assert.throws(
129    () =>
130      resolveSeedreamSize("doubao-seededit-3-0-i2i-250628", makeArgs({ size: "1024x1024" })),
131    /no longer supported/,
132  );
133});
134 
135test("Seedream reference-image support is model-specific", () => {
136  assert.doesNotThrow(() =>
137    validateArgs(
138      "doubao-seedream-5-0-260128",
139      makeArgs({ referenceImages: ["a.png", "b.png"] }),
140    ),
141  );
142 
143  assert.throws(
144    () =>
145      validateArgs(
146        "doubao-seedream-3-0-t2i-250415",
147        makeArgs({ referenceImages: ["a.png"] }),
148      ),
149    /does not support reference images/,
150  );
151 
152  assert.throws(
153    () =>
154      validateArgs(
155        "doubao-seededit-3-0-i2i-250628",
156        makeArgs(),
157      ),
158    /no longer supported/,
159  );
160 
161  assert.throws(
162    () =>
163      validateArgs(
164        "ep-20260315171508-t8br2",
165        makeArgs({ referenceImages: ["a.png"] }),
166      ),
167    /require a known model ID/,
168  );
169});
170 
171test("Seedream image input encodes local references as data URLs", async (t) => {
172  const refOne = await makeTempPng(t, "one.png");
173  const refTwo = await makeTempPng(t, "two.png");
174 
175  const single = await buildImageInput("doubao-seedream-4-5-251128", [refOne]);
176  assert.match(String(single), /^data:image\/png;base64,/);
177 
178  const multiple = await buildImageInput("doubao-seedream-5-0-260128", [refOne, refTwo]);
179  assert.ok(Array.isArray(multiple));
180  assert.equal(multiple.length, 2);
181});
182 
183test("Seedream generateImage posts the documented response_format and downloads the returned URL", async (t) => {
184  useEnv(t, { ARK_API_KEY: "test-key", SEEDREAM_BASE_URL: null });
185 
186  const originalFetch = globalThis.fetch;
187  t.after(() => {
188    globalThis.fetch = originalFetch;
189  });
190 
191  const calls: Array<{
192    input: string;
193    init?: RequestInit;
194  }> = [];
195 
196  globalThis.fetch = async (input, init) => {
197    calls.push({
198      input: String(input),
199      init,
200    });
201 
202    if (calls.length === 1) {
203      return Response.json({
204        model: "doubao-seedream-4-5-251128",
205        created: 1740000000,
206        data: [
207          {
208            url: "https://example.com/generated-image",
209            size: "2048x2048",
210          },
211        ],
212        usage: {
213          generated_images: 1,
214          output_tokens: 1,
215          total_tokens: 1,
216        },
217      });
218    }
219 
220    return new Response(Uint8Array.from([7, 8, 9]), {
221      status: 200,
222      headers: { "Content-Type": "image/jpeg" },
223    });
224  };
225 
226  const image = await generateImage(
227    "A robot illustrator",
228    "doubao-seedream-4-5-251128",
229    makeArgs(),
230  );
231 
232  assert.deepEqual([...image], [7, 8, 9]);
233  assert.equal(calls.length, 2);
234  assert.equal(
235    calls[0]?.input,
236    "https://ark.cn-beijing.volces.com/api/v3/images/generations",
237  );
238 
239  const requestBody = JSON.parse(String(calls[0]?.init?.body)) as Record<string, unknown>;
240  assert.equal(requestBody.model, "doubao-seedream-4-5-251128");
241  assert.equal(requestBody.size, "2K");
242  assert.equal(requestBody.response_format, "url");
243  assert.ok(!("output_format" in requestBody));
244  assert.equal(calls[1]?.input, "https://example.com/generated-image");
245});
246