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