mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
test: major coverage expansion — 18 new test files, ~830 new tests
Unit tests: 1354 → 1781 (+427) - 11 new AI bridge module tests (packages/ai/ from 2/13 → 13/13 files) - files-page-store (0% → full), pdf-to-image-store, features-store expanded - saturation and edit-metadata image-engine operations - analytics route, features route, web analytics lib, api-extended Integration tests: ~2070 → 2320 (+250) - 31 integration files expanded with branch-coverage-targeted tests - progress.ts SSE endpoints (28% → comprehensive, +18 tests) - gif-tools all modes (+18), pdf-to-image format variants (+13) - Cross-format matrix expanded to 17 tools × 17 formats (467 tests) - Adversarial: concurrent, memory pressure, unicode filenames, pipeline limits E2E-Docker: +1020 lines across 6 spec files - Info, colors, sharpening, base64, QR read, JXL/ICO/SVG formats - Strip-metadata, image-enhancement, content-aware-resize expanded - Batch pipelines, multi-format batches, HEIC input coverage
This commit is contained in:
SnapOtter
@@ -0,0 +1,298 @@
|
||||
import { readFile, unlink, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
unlink: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { removeBackground } from "../../../packages/ai/src/background-removal.js";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-image-data");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-output";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(unlink).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: true });
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("removeBackground", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls remove_bg.py with input path, output path, and options JSON", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"remove_bg.py",
|
||||
expect.arrayContaining([
|
||||
expect.stringContaining("rembg_in_"),
|
||||
expect.stringContaining("rembg_out_"),
|
||||
"{}",
|
||||
]),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes model option into the args JSON", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "u2net_human_seg" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ model: "u2net_human_seg" });
|
||||
});
|
||||
|
||||
it("serializes backgroundColor option into the args JSON", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, { backgroundColor: "#FF0000" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ backgroundColor: "#FF0000" });
|
||||
});
|
||||
|
||||
it("serializes both model and backgroundColor together", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, {
|
||||
model: "isnet-general-use",
|
||||
backgroundColor: "transparent",
|
||||
});
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({
|
||||
model: "isnet-general-use",
|
||||
backgroundColor: "transparent",
|
||||
});
|
||||
});
|
||||
|
||||
it("converts input to PNG via sharp before writing to disk", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
expect.stringContaining("rembg_in_"),
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
|
||||
it("uses unique UUID in temp file names to prevent collisions", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
const call1Args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
const call2Args = vi.mocked(runPythonWithProgress).mock.calls[1][1];
|
||||
// The UUID portions should differ
|
||||
expect(call1Args[0]).not.toBe(call2Args[0]);
|
||||
});
|
||||
|
||||
it("writes input to system tmpdir, output to outputDir", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
// Input path in tmpdir
|
||||
expect(args[0]).toMatch(/rembg_in_/);
|
||||
// Output path in outputDir
|
||||
expect(args[1]).toMatch(/^\/tmp\/test-output\/rembg_out_/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("reads the output file and returns its buffer on success", async () => {
|
||||
const outputBuf = Buffer.from("transparent-image");
|
||||
vi.mocked(readFile).mockResolvedValueOnce(outputBuf);
|
||||
|
||||
const result = await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result).toBe(outputBuf);
|
||||
});
|
||||
|
||||
it("passes stdout through parseStdoutJson", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true, "extra": "data"}',
|
||||
stderr: "",
|
||||
});
|
||||
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(parseStdoutJson).toHaveBeenCalledWith('{"success": true, "extra": "data"}');
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws when Python returns success: false with custom error", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "Model u2net_cloth not available",
|
||||
});
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Model u2net_cloth not available",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false and no error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Background removal failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates bridge timeout errors", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Python script timed out",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates parseStdoutJson errors", async () => {
|
||||
vi.mocked(parseStdoutJson).mockImplementation(() => {
|
||||
throw new Error("No JSON response from Python script");
|
||||
});
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"No JSON response from Python script",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates sharp conversion errors", async () => {
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockRejectedValue(new Error("Invalid image")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("Invalid image");
|
||||
});
|
||||
});
|
||||
|
||||
describe("timeout calculation", () => {
|
||||
it("uses 300000ms base timeout for non-birefnet models", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "u2net" });
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBe(300000);
|
||||
});
|
||||
|
||||
it("uses 600000ms base timeout for birefnet models", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "birefnet-general" });
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBeGreaterThanOrEqual(600000);
|
||||
});
|
||||
|
||||
it("uses 600000ms base timeout for birefnet-massive model", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "birefnet-massive" });
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBeGreaterThanOrEqual(600000);
|
||||
});
|
||||
|
||||
it("scales timeout with megapixels for large images", async () => {
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 6000, height: 4000 }),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
// 24 MP * 30 * 1000 = 720000 > 300000
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBe(720000);
|
||||
});
|
||||
|
||||
it("uses default 300000ms base when model is not specified", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBe(300000);
|
||||
});
|
||||
});
|
||||
|
||||
describe("temp file cleanup", () => {
|
||||
it("cleans up both input and output temp files on success", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(unlink).toHaveBeenCalledTimes(2);
|
||||
expect(unlink).toHaveBeenCalledWith(expect.stringContaining("rembg_in_"));
|
||||
expect(unlink).toHaveBeenCalledWith(expect.stringContaining("rembg_out_"));
|
||||
});
|
||||
|
||||
it("cleans up temp files when Python returns failure", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow();
|
||||
expect(unlink).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("cleans up temp files when bridge rejects", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("crash"));
|
||||
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow();
|
||||
expect(unlink).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("does not throw if unlink fails (swallows cleanup errors)", async () => {
|
||||
vi.mocked(unlink).mockRejectedValue(new Error("ENOENT"));
|
||||
|
||||
// Should not throw -- unlink errors are caught
|
||||
await expect(removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR)).resolves.toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress callback through to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"remove_bg.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
|
||||
it("passes undefined onProgress when not provided", async () => {
|
||||
await removeBackground(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.onProgress).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -618,3 +618,235 @@ describe("bridge - parseStdoutJson edge cases", () => {
|
||||
expect(result).toEqual({ success: true, device: "cpu" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("bridge - getDispatcherStatus", () => {
|
||||
let getDispatcherStatus: typeof import("../../../packages/ai/src/bridge.js").getDispatcherStatus;
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.resetModules();
|
||||
const mod = await import("../../../packages/ai/src/bridge.js");
|
||||
getDispatcherStatus = mod.getDispatcherStatus;
|
||||
});
|
||||
|
||||
it("returns initial state with no dispatcher running", () => {
|
||||
const status = getDispatcherStatus();
|
||||
expect(status).toEqual({
|
||||
running: false,
|
||||
ready: false,
|
||||
failed: false,
|
||||
gpu: false,
|
||||
pid: null,
|
||||
consecutiveCrashes: 0,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("bridge - dispatcher lifecycle via runPythonWithProgress", () => {
|
||||
let runPythonWithProgress: typeof import("../../../packages/ai/src/bridge.js").runPythonWithProgress;
|
||||
let getDispatcherStatus: typeof import("../../../packages/ai/src/bridge.js").getDispatcherStatus;
|
||||
let shutdownDispatcher: typeof import("../../../packages/ai/src/bridge.js").shutdownDispatcher;
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.resetModules();
|
||||
vi.mocked(spawn).mockReset();
|
||||
|
||||
const mod = await import("../../../packages/ai/src/bridge.js");
|
||||
runPythonWithProgress = mod.runPythonWithProgress;
|
||||
getDispatcherStatus = mod.getDispatcherStatus;
|
||||
shutdownDispatcher = mod.shutdownDispatcher;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it("falls back to per-request spawn when dispatcher ENOENT marks it failed", async () => {
|
||||
const mockDispatcher = createMockProcess();
|
||||
const mockPerRequest = createMockProcess();
|
||||
let callCount = 0;
|
||||
|
||||
vi.mocked(spawn).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return mockDispatcher.process;
|
||||
return mockPerRequest.process;
|
||||
});
|
||||
|
||||
const promise = runPythonWithProgress("test.py", ["arg1"]);
|
||||
|
||||
// Dispatcher fails with ENOENT => permanently failed
|
||||
const enoent = new Error("spawn ENOENT") as NodeJS.ErrnoException;
|
||||
enoent.code = "ENOENT";
|
||||
mockDispatcher.emitEvent("error", enoent);
|
||||
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
|
||||
// Per-request spawn succeeds
|
||||
mockPerRequest.stdout.emit("data", Buffer.from('{"ok": true}\n'));
|
||||
mockPerRequest.emitEvent("close", 0, null);
|
||||
|
||||
const result = await promise;
|
||||
expect(result.stdout).toContain('{"ok": true}');
|
||||
});
|
||||
|
||||
it("reports failed status when dispatcher ENOENT occurs", async () => {
|
||||
const mockDispatcher = createMockProcess();
|
||||
const mockPerRequest = createMockProcess();
|
||||
let callCount = 0;
|
||||
|
||||
vi.mocked(spawn).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return mockDispatcher.process;
|
||||
return mockPerRequest.process;
|
||||
});
|
||||
|
||||
const promise = runPythonWithProgress("test.py", []);
|
||||
|
||||
const enoent = new Error("spawn ENOENT") as NodeJS.ErrnoException;
|
||||
enoent.code = "ENOENT";
|
||||
mockDispatcher.emitEvent("error", enoent);
|
||||
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
|
||||
// Finish the per-request
|
||||
mockPerRequest.stdout.emit("data", Buffer.from('{"ok": true}\n'));
|
||||
mockPerRequest.emitEvent("close", 0, null);
|
||||
await promise;
|
||||
|
||||
const status = getDispatcherStatus();
|
||||
expect(status.failed).toBe(true);
|
||||
expect(status.running).toBe(false);
|
||||
});
|
||||
|
||||
it("graceful shutdown does not throw when dispatcher already exited", async () => {
|
||||
const mockDispatcher = createMockProcess();
|
||||
const mockPerReq = createMockProcess();
|
||||
let callCount = 0;
|
||||
|
||||
vi.mocked(spawn).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return mockDispatcher.process;
|
||||
return mockPerReq.process;
|
||||
});
|
||||
|
||||
const promise = runPythonWithProgress("test.py", []);
|
||||
|
||||
// Dispatcher closes (crash) -- sets dispatcher = null internally
|
||||
mockDispatcher.emitEvent("close", 1, null);
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
|
||||
// shutdownDispatcher should not throw even when no dispatcher is running
|
||||
expect(() => shutdownDispatcher()).not.toThrow();
|
||||
|
||||
// Finish the per-request fallback
|
||||
mockPerReq.stdout.emit("data", Buffer.from('{"ok": true}\n'));
|
||||
mockPerReq.emitEvent("close", 0, null);
|
||||
|
||||
await promise;
|
||||
});
|
||||
|
||||
it("shutdown is idempotent when called multiple times", () => {
|
||||
expect(() => {
|
||||
shutdownDispatcher();
|
||||
shutdownDispatcher();
|
||||
shutdownDispatcher();
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it("concurrent requests to per-request fallback both resolve", async () => {
|
||||
// Dispatcher fails immediately, so both requests go to per-request path
|
||||
const mockDispatcher = createMockProcess();
|
||||
const mockReq1 = createMockProcess();
|
||||
const mockReq2 = createMockProcess();
|
||||
let callCount = 0;
|
||||
|
||||
vi.mocked(spawn).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return mockDispatcher.process;
|
||||
if (callCount === 2) return mockReq1.process;
|
||||
return mockReq2.process;
|
||||
});
|
||||
|
||||
// Start first request
|
||||
const promise1 = runPythonWithProgress("tool1.py", ["a"]);
|
||||
|
||||
// Kill dispatcher
|
||||
const enoent = new Error("spawn ENOENT") as NodeJS.ErrnoException;
|
||||
enoent.code = "ENOENT";
|
||||
mockDispatcher.emitEvent("error", enoent);
|
||||
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
|
||||
// Start second request (dispatcher is now permanently failed)
|
||||
const promise2 = runPythonWithProgress("tool2.py", ["b"]);
|
||||
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
|
||||
// Complete both per-request processes
|
||||
mockReq1.stdout.emit("data", Buffer.from('{"result": "one"}\n'));
|
||||
mockReq1.emitEvent("close", 0, null);
|
||||
|
||||
mockReq2.stdout.emit("data", Buffer.from('{"result": "two"}\n'));
|
||||
mockReq2.emitEvent("close", 0, null);
|
||||
|
||||
const [r1, r2] = await Promise.all([promise1, promise2]);
|
||||
expect(r1.stdout).toContain("one");
|
||||
expect(r2.stdout).toContain("two");
|
||||
});
|
||||
|
||||
it("timeout rejects the promise without affecting other requests", async () => {
|
||||
vi.useFakeTimers();
|
||||
const mockDispatcher = createMockProcess();
|
||||
const mockReq = createMockProcess();
|
||||
let callCount = 0;
|
||||
|
||||
vi.mocked(spawn).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return mockDispatcher.process;
|
||||
return mockReq.process;
|
||||
});
|
||||
|
||||
const promise = runPythonWithProgress("slow.py", [], { timeout: 2000 });
|
||||
|
||||
// Dispatcher ENOENT => per-request fallback
|
||||
const enoent = new Error("spawn ENOENT") as NodeJS.ErrnoException;
|
||||
enoent.code = "ENOENT";
|
||||
mockDispatcher.emitEvent("error", enoent);
|
||||
|
||||
// Advance past timeout
|
||||
vi.advanceTimersByTime(3000);
|
||||
|
||||
// Process gets killed, close fires
|
||||
mockReq.emitEvent("close", null, "SIGTERM");
|
||||
|
||||
await expect(promise).rejects.toThrow("Python script timed out");
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("handles dispatcher crash followed by successful per-request retry", async () => {
|
||||
const mockDispatcher = createMockProcess();
|
||||
const mockPerReq = createMockProcess();
|
||||
let callCount = 0;
|
||||
|
||||
vi.mocked(spawn).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return mockDispatcher.process;
|
||||
return mockPerReq.process;
|
||||
});
|
||||
|
||||
const promise = runPythonWithProgress("test.py", []);
|
||||
|
||||
// Dispatcher crashes with a non-ENOENT error
|
||||
const err = new Error("spawn failed");
|
||||
(err as NodeJS.ErrnoException).code = "EACCES";
|
||||
mockDispatcher.emitEvent("error", err);
|
||||
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
|
||||
// Per-request succeeds
|
||||
mockPerReq.stdout.emit("data", Buffer.from('{"success": true}\n'));
|
||||
mockPerReq.emitEvent("close", 0, null);
|
||||
|
||||
const result = await promise;
|
||||
expect(result.stdout).toContain("success");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { colorize } from "../../../packages/ai/src/colorization.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-bw-image");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-colorize";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true, "width": 800, "height": 600, "method": "deoldify"}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
method: "deoldify",
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("colorize", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls colorize.py with input path, output path, and options JSON", async () => {
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"colorize.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_colorize.png`, `${FAKE_OUTPUT_DIR}/output_colorize.png`, "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes intensity option", async () => {
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR, { intensity: 0.5 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ intensity: 0.5 });
|
||||
});
|
||||
|
||||
it("serializes model option", async () => {
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "eccv16" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ model: "eccv16" });
|
||||
});
|
||||
|
||||
it("serializes both intensity and model together", async () => {
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR, { intensity: 1.0, model: "siggraph17" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ intensity: 1.0, model: "siggraph17" });
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
`${FAKE_OUTPUT_DIR}/input_colorize.png`,
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns ColorizeResult with buffer, width, height, and method", async () => {
|
||||
const result = await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
buffer: expect.any(Buffer),
|
||||
width: 800,
|
||||
height: 600,
|
||||
method: "deoldify",
|
||||
});
|
||||
});
|
||||
|
||||
it("reads from default output path when output_path not in response", async () => {
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_colorize.png`);
|
||||
});
|
||||
|
||||
it("reads from alternate output_path when provided by Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
output_path: "/tmp/alternate-colorized.png",
|
||||
});
|
||||
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(readFile).toHaveBeenCalledWith("/tmp/alternate-colorized.png");
|
||||
});
|
||||
|
||||
it("defaults method to 'unknown' when not provided by Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.method).toBe("unknown");
|
||||
});
|
||||
|
||||
it("preserves width and height from Python response", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
method: "eccv16",
|
||||
});
|
||||
|
||||
const result = await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.width).toBe(1920);
|
||||
expect(result.height).toBe(1080);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws when Python returns success: false with custom error", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "Input is already a color image",
|
||||
});
|
||||
|
||||
await expect(colorize(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Input is already a color image",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false and no error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(colorize(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("Colorization failed");
|
||||
});
|
||||
|
||||
it("propagates bridge rejection", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process killed (out of memory)"),
|
||||
);
|
||||
|
||||
await expect(colorize(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("out of memory");
|
||||
});
|
||||
|
||||
it("propagates parseStdoutJson errors", async () => {
|
||||
vi.mocked(parseStdoutJson).mockImplementation(() => {
|
||||
throw new SyntaxError("Unexpected token");
|
||||
});
|
||||
|
||||
await expect(colorize(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("Unexpected token");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress callback to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await colorize(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"colorize.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,299 @@
|
||||
import { readFile, unlink, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
unlink: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { blurFaces, detectFaces } from "../../../packages/ai/src/face-detection.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-image-data");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-faces";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(unlink).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true, "facesDetected": 0}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 0,
|
||||
faces: [],
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("blurFaces", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls detect_faces.py with correct file paths", async () => {
|
||||
await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"detect_faces.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_faces.png`, `${FAKE_OUTPUT_DIR}/output_faces.png`, "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes blurRadius option", async () => {
|
||||
await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { blurRadius: 30 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ blurRadius: 30 });
|
||||
});
|
||||
|
||||
it("serializes sensitivity option", async () => {
|
||||
await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { sensitivity: 0.3 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ sensitivity: 0.3 });
|
||||
});
|
||||
|
||||
it("serializes both blurRadius and sensitivity", async () => {
|
||||
await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { blurRadius: 25, sensitivity: 0.7 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ blurRadius: 25, sensitivity: 0.7 });
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns BlurFacesResult with multiple face regions", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 3,
|
||||
faces: [
|
||||
{ x: 10, y: 20, w: 50, h: 60 },
|
||||
{ x: 100, y: 120, w: 55, h: 65 },
|
||||
{ x: 200, y: 220, w: 45, h: 55 },
|
||||
],
|
||||
});
|
||||
|
||||
const result = await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result.facesDetected).toBe(3);
|
||||
expect(result.faces).toHaveLength(3);
|
||||
expect(result.buffer).toBeInstanceOf(Buffer);
|
||||
});
|
||||
|
||||
it("returns empty faces array when none detected", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 0,
|
||||
});
|
||||
|
||||
const result = await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.facesDetected).toBe(0);
|
||||
expect(result.faces).toEqual([]);
|
||||
});
|
||||
|
||||
it("defaults faces to empty array when field absent from response", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 0,
|
||||
});
|
||||
|
||||
const result = await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.faces).toEqual([]);
|
||||
});
|
||||
|
||||
it("reads the output file for the blurred image", async () => {
|
||||
const blurredBuf = Buffer.from("blurred-faces-output");
|
||||
vi.mocked(readFile).mockResolvedValueOnce(blurredBuf);
|
||||
|
||||
const result = await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.buffer).toBe(blurredBuf);
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_faces.png`);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "MediaPipe initialization failed",
|
||||
});
|
||||
|
||||
await expect(blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"MediaPipe initialization failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("Face detection failed");
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await blurFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"detect_faces.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("detectFaces", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls detect_faces.py with detectOnly: true", async () => {
|
||||
await detectFaces(FAKE_INPUT);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
const optionsArg = JSON.parse(args[2]);
|
||||
expect(optionsArg.detectOnly).toBe(true);
|
||||
});
|
||||
|
||||
it("passes 'unused' as the output path argument", async () => {
|
||||
await detectFaces(FAKE_INPUT);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(args[1]).toBe("unused");
|
||||
});
|
||||
|
||||
it("merges user sensitivity with detectOnly flag", async () => {
|
||||
await detectFaces(FAKE_INPUT, { sensitivity: 0.2 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
const parsed = JSON.parse(args[2]);
|
||||
expect(parsed).toEqual({ sensitivity: 0.2, detectOnly: true });
|
||||
});
|
||||
|
||||
it("writes input to tmpdir", async () => {
|
||||
await detectFaces(FAKE_INPUT);
|
||||
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
expect.stringContaining("detect_faces_"),
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns DetectFacesResult without a buffer property", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 2,
|
||||
faces: [
|
||||
{ x: 10, y: 20, w: 50, h: 60 },
|
||||
{ x: 100, y: 120, w: 55, h: 65 },
|
||||
],
|
||||
});
|
||||
|
||||
const result = await detectFaces(FAKE_INPUT);
|
||||
|
||||
expect(result.facesDetected).toBe(2);
|
||||
expect(result.faces).toHaveLength(2);
|
||||
expect(result).not.toHaveProperty("buffer");
|
||||
});
|
||||
|
||||
it("defaults faces to empty array", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 0,
|
||||
});
|
||||
|
||||
const result = await detectFaces(FAKE_INPUT);
|
||||
expect(result.faces).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("temp file cleanup", () => {
|
||||
it("cleans up temp input file after success", async () => {
|
||||
await detectFaces(FAKE_INPUT);
|
||||
expect(unlink).toHaveBeenCalledWith(expect.stringContaining("detect_faces_"));
|
||||
});
|
||||
|
||||
it("cleans up temp input file after failure", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(detectFaces(FAKE_INPUT)).rejects.toThrow();
|
||||
expect(unlink).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("cleans up temp file when bridge rejects", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("crash"));
|
||||
|
||||
await expect(detectFaces(FAKE_INPUT)).rejects.toThrow();
|
||||
expect(unlink).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws fallback error", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(detectFaces(FAKE_INPUT)).rejects.toThrow("Face detection failed");
|
||||
});
|
||||
|
||||
it("propagates segfault error", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process crashed (segmentation fault)"),
|
||||
);
|
||||
|
||||
await expect(detectFaces(FAKE_INPUT)).rejects.toThrow("segmentation fault");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await detectFaces(FAKE_INPUT, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"detect_faces.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,229 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { enhanceFaces } from "../../../packages/ai/src/face-enhancement.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-image-data");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-enhance";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true, "facesDetected": 1, "faces": [], "model": "gfpgan"}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 1,
|
||||
faces: [{ x: 10, y: 20, w: 80, h: 90 }],
|
||||
model: "gfpgan",
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("enhanceFaces", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls enhance_faces.py with correct file paths", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"enhance_faces.py",
|
||||
[
|
||||
`${FAKE_OUTPUT_DIR}/input_enhance_faces.png`,
|
||||
`${FAKE_OUTPUT_DIR}/output_enhance_faces.png`,
|
||||
"{}",
|
||||
],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes model option", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "codeformer" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ model: "codeformer" });
|
||||
});
|
||||
|
||||
it("serializes strength option", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { strength: 0.7 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ strength: 0.7 });
|
||||
});
|
||||
|
||||
it("serializes onlyCenterFace option", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { onlyCenterFace: true });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ onlyCenterFace: true });
|
||||
});
|
||||
|
||||
it("serializes sensitivity option", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, { sensitivity: 0.4 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ sensitivity: 0.4 });
|
||||
});
|
||||
|
||||
it("serializes all options together", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, {
|
||||
model: "auto",
|
||||
strength: 0.5,
|
||||
onlyCenterFace: false,
|
||||
sensitivity: 0.6,
|
||||
});
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({
|
||||
model: "auto",
|
||||
strength: 0.5,
|
||||
onlyCenterFace: false,
|
||||
sensitivity: 0.6,
|
||||
});
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
`${FAKE_OUTPUT_DIR}/input_enhance_faces.png`,
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns EnhanceFacesResult with all fields", async () => {
|
||||
const result = await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
buffer: expect.any(Buffer),
|
||||
facesDetected: 1,
|
||||
faces: [{ x: 10, y: 20, w: 80, h: 90 }],
|
||||
model: "gfpgan",
|
||||
});
|
||||
});
|
||||
|
||||
it("reads the enhanced output file", async () => {
|
||||
const enhancedBuf = Buffer.from("enhanced-faces");
|
||||
vi.mocked(readFile).mockResolvedValueOnce(enhancedBuf);
|
||||
|
||||
const result = await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.buffer).toBe(enhancedBuf);
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_enhance_faces.png`);
|
||||
});
|
||||
|
||||
it("defaults model to 'unknown' when absent from response", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 1,
|
||||
faces: [{ x: 0, y: 0, w: 50, h: 50 }],
|
||||
});
|
||||
|
||||
const result = await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.model).toBe("unknown");
|
||||
});
|
||||
|
||||
it("defaults faces to empty array when absent from response", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 0,
|
||||
});
|
||||
|
||||
const result = await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.faces).toEqual([]);
|
||||
});
|
||||
|
||||
it("returns multiple face regions", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 3,
|
||||
faces: [
|
||||
{ x: 10, y: 10, w: 50, h: 50 },
|
||||
{ x: 100, y: 100, w: 60, h: 60 },
|
||||
{ x: 200, y: 50, w: 40, h: 40 },
|
||||
],
|
||||
model: "codeformer",
|
||||
});
|
||||
|
||||
const result = await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.facesDetected).toBe(3);
|
||||
expect(result.faces).toHaveLength(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "GFPGAN weights not found",
|
||||
});
|
||||
|
||||
await expect(enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"GFPGAN weights not found",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Face enhancement failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates OOM errors from bridge", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process killed (out of memory)"),
|
||||
);
|
||||
|
||||
await expect(enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("out of memory");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await enhanceFaces(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"enhance_faces.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,266 @@
|
||||
import { unlink, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
unlink: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { detectFaceLandmarks } from "../../../packages/ai/src/face-landmarks.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-image-data");
|
||||
|
||||
const FULL_LANDMARKS = {
|
||||
leftEye: { x: 100, y: 150 },
|
||||
rightEye: { x: 200, y: 150 },
|
||||
eyeCenter: { x: 150, y: 150 },
|
||||
chin: { x: 150, y: 300 },
|
||||
forehead: { x: 150, y: 80 },
|
||||
crown: { x: 150, y: 50 },
|
||||
nose: { x: 150, y: 200 },
|
||||
faceCenterX: 150,
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(unlink).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
faceDetected: true,
|
||||
landmarks: FULL_LANDMARKS,
|
||||
imageWidth: 800,
|
||||
imageHeight: 600,
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("detectFaceLandmarks", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls face_landmarks.py with input path, 'unused', and '{}'", async () => {
|
||||
await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"face_landmarks.py",
|
||||
[expect.stringContaining("face_landmarks_"), "unused", "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("writes input buffer directly without sharp conversion", async () => {
|
||||
await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
// face-landmarks.ts does NOT use sharp -- it writes inputBuffer directly
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
expect.stringContaining("face_landmarks_"),
|
||||
FAKE_INPUT,
|
||||
);
|
||||
});
|
||||
|
||||
it("writes to system tmpdir", async () => {
|
||||
await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
const writePath = vi.mocked(writeFile).mock.calls[0][0] as string;
|
||||
expect(writePath).toContain("face_landmarks_");
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns all landmark points when face is detected", async () => {
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
expect(result.faceDetected).toBe(true);
|
||||
expect(result.landmarks).toEqual(FULL_LANDMARKS);
|
||||
});
|
||||
|
||||
it("returns individual landmark points correctly", async () => {
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
expect(result.landmarks!.leftEye).toEqual({ x: 100, y: 150 });
|
||||
expect(result.landmarks!.rightEye).toEqual({ x: 200, y: 150 });
|
||||
expect(result.landmarks!.eyeCenter).toEqual({ x: 150, y: 150 });
|
||||
expect(result.landmarks!.chin).toEqual({ x: 150, y: 300 });
|
||||
expect(result.landmarks!.forehead).toEqual({ x: 150, y: 80 });
|
||||
expect(result.landmarks!.crown).toEqual({ x: 150, y: 50 });
|
||||
expect(result.landmarks!.nose).toEqual({ x: 150, y: 200 });
|
||||
expect(result.landmarks!.faceCenterX).toBe(150);
|
||||
});
|
||||
|
||||
it("returns imageWidth and imageHeight from response", async () => {
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
expect(result.imageWidth).toBe(800);
|
||||
expect(result.imageHeight).toBe(600);
|
||||
});
|
||||
|
||||
it("returns null landmarks when no face detected", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
faceDetected: false,
|
||||
imageWidth: 1024,
|
||||
imageHeight: 768,
|
||||
});
|
||||
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
expect(result.faceDetected).toBe(false);
|
||||
expect(result.landmarks).toBeNull();
|
||||
});
|
||||
|
||||
it("defaults landmarks to null when field is absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
faceDetected: true,
|
||||
imageWidth: 800,
|
||||
imageHeight: 600,
|
||||
});
|
||||
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
expect(result.landmarks).toBeNull();
|
||||
});
|
||||
|
||||
it("defaults imageWidth to 0 when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
faceDetected: false,
|
||||
});
|
||||
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
expect(result.imageWidth).toBe(0);
|
||||
});
|
||||
|
||||
it("defaults imageHeight to 0 when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
faceDetected: false,
|
||||
});
|
||||
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
expect(result.imageHeight).toBe(0);
|
||||
});
|
||||
|
||||
it("returns both dimensions as 0 when both absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
faceDetected: true,
|
||||
landmarks: FULL_LANDMARKS,
|
||||
});
|
||||
|
||||
const result = await detectFaceLandmarks(FAKE_INPUT);
|
||||
expect(result.imageWidth).toBe(0);
|
||||
expect(result.imageHeight).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "MediaPipe model not found",
|
||||
});
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow("MediaPipe model not found");
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow(
|
||||
"Face landmark detection failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow("timed out");
|
||||
});
|
||||
|
||||
it("propagates OOM errors", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process killed (out of memory)"),
|
||||
);
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow("out of memory");
|
||||
});
|
||||
|
||||
it("propagates parseStdoutJson errors", async () => {
|
||||
vi.mocked(parseStdoutJson).mockImplementation(() => {
|
||||
throw new Error("No JSON response from Python script");
|
||||
});
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow(
|
||||
"No JSON response from Python script",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("temp file cleanup", () => {
|
||||
it("cleans up temp input file after success", async () => {
|
||||
await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
expect(unlink).toHaveBeenCalledWith(expect.stringContaining("face_landmarks_"));
|
||||
});
|
||||
|
||||
it("cleans up temp input file when Python returns failure", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow();
|
||||
expect(unlink).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("cleans up temp input file when bridge rejects", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("crash"));
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).rejects.toThrow();
|
||||
expect(unlink).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not throw if unlink fails", async () => {
|
||||
vi.mocked(unlink).mockRejectedValue(new Error("ENOENT"));
|
||||
|
||||
await expect(detectFaceLandmarks(FAKE_INPUT)).resolves.toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await detectFaceLandmarks(FAKE_INPUT, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"face_landmarks.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
|
||||
it("omits onProgress when not provided", async () => {
|
||||
await detectFaceLandmarks(FAKE_INPUT);
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.onProgress).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,202 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { inpaint } from "../../../packages/ai/src/inpainting.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-image-data");
|
||||
const FAKE_MASK = Buffer.from("fake-mask-data");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-inpaint";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: true });
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("inpaint", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls inpaint.py with input, mask, and output paths", async () => {
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"inpaint.py",
|
||||
[
|
||||
`${FAKE_OUTPUT_DIR}/input_inpaint.png`,
|
||||
`${FAKE_OUTPUT_DIR}/mask_inpaint.png`,
|
||||
`${FAKE_OUTPUT_DIR}/output_inpaint.png`,
|
||||
],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("converts both input and mask to PNG via sharp", async () => {
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
|
||||
// sharp called twice: once for input, once for mask
|
||||
expect(sharp).toHaveBeenCalledTimes(2);
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_MASK);
|
||||
});
|
||||
|
||||
it("writes both input and mask files to disk", async () => {
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(writeFile).toHaveBeenCalledTimes(2);
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
`${FAKE_OUTPUT_DIR}/input_inpaint.png`,
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
`${FAKE_OUTPUT_DIR}/mask_inpaint.png`,
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
|
||||
it("does not pass any options argument (only 3 args)", async () => {
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(args).toHaveLength(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns output buffer on success", async () => {
|
||||
const inpaintedBuf = Buffer.from("inpainted-result");
|
||||
vi.mocked(readFile).mockResolvedValueOnce(inpaintedBuf);
|
||||
|
||||
const result = await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
expect(result).toBe(inpaintedBuf);
|
||||
});
|
||||
|
||||
it("reads from the correct output path", async () => {
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_inpaint.png`);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "Mask dimensions do not match input",
|
||||
});
|
||||
|
||||
await expect(inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Mask dimensions do not match input",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Inpainting failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
|
||||
it("propagates bridge OOM", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process killed (out of memory)"),
|
||||
);
|
||||
|
||||
await expect(inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"out of memory",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates parseStdoutJson errors", async () => {
|
||||
vi.mocked(parseStdoutJson).mockImplementation(() => {
|
||||
throw new Error("No JSON response from Python script");
|
||||
});
|
||||
|
||||
await expect(inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"No JSON response from Python script",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates sharp conversion errors on input", async () => {
|
||||
let callCount = 0;
|
||||
vi.mocked(sharp).mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) {
|
||||
return {
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockRejectedValue(new Error("Corrupt input image")),
|
||||
} as unknown as ReturnType<typeof sharp>;
|
||||
}
|
||||
return {
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
} as unknown as ReturnType<typeof sharp>;
|
||||
});
|
||||
|
||||
await expect(inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Corrupt input image",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"inpaint.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
|
||||
it("omits onProgress when not provided", async () => {
|
||||
await inpaint(FAKE_INPUT, FAKE_MASK, FAKE_OUTPUT_DIR);
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.onProgress).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,242 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { noiseRemoval } from "../../../packages/ai/src/noise-removal.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-noisy-image");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-denoise";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
format: "png",
|
||||
tier: "balanced",
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("noiseRemoval", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls noise_removal.py with correct file paths", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"noise_removal.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_denoise.png`, `${FAKE_OUTPUT_DIR}/output_denoise.png`, "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes tier option", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, { tier: "aggressive" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ tier: "aggressive" });
|
||||
});
|
||||
|
||||
it("serializes strength option", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, { strength: 0.8 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ strength: 0.8 });
|
||||
});
|
||||
|
||||
it("serializes detailPreservation option", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, { detailPreservation: 0.6 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ detailPreservation: 0.6 });
|
||||
});
|
||||
|
||||
it("serializes colorNoise option", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, { colorNoise: 0.4 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ colorNoise: 0.4 });
|
||||
});
|
||||
|
||||
it("serializes format and quality options", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, { format: "webp", quality: 90 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ format: "webp", quality: 90 });
|
||||
});
|
||||
|
||||
it("serializes all options together", async () => {
|
||||
const allOptions = {
|
||||
tier: "aggressive",
|
||||
strength: 0.9,
|
||||
detailPreservation: 0.5,
|
||||
colorNoise: 0.3,
|
||||
format: "jpeg",
|
||||
quality: 85,
|
||||
};
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, allOptions);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual(allOptions);
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns NoiseRemovalResult with all fields", async () => {
|
||||
const result = await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
buffer: expect.any(Buffer),
|
||||
width: 800,
|
||||
height: 600,
|
||||
format: "png",
|
||||
tier: "balanced",
|
||||
});
|
||||
});
|
||||
|
||||
it("reads from default output path when output_path not in response", async () => {
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_denoise.png`);
|
||||
});
|
||||
|
||||
it("reads from alternate output_path when provided", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
output_path: "/tmp/alt-denoise.webp",
|
||||
});
|
||||
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(readFile).toHaveBeenCalledWith("/tmp/alt-denoise.webp");
|
||||
});
|
||||
|
||||
it("defaults format to 'png' when absent from response", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.format).toBe("png");
|
||||
});
|
||||
|
||||
it("defaults tier from Python response when present", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
tier: "gentle",
|
||||
});
|
||||
|
||||
const result = await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.tier).toBe("gentle");
|
||||
});
|
||||
|
||||
it("falls back to options tier when Python omits it", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, { tier: "aggressive" });
|
||||
expect(result.tier).toBe("aggressive");
|
||||
});
|
||||
|
||||
it("falls back to 'balanced' when both Python and options omit tier", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.tier).toBe("balanced");
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "NAFNet model loading failed",
|
||||
});
|
||||
|
||||
await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"NAFNet model loading failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Noise removal failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"noise_removal.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,261 @@
|
||||
import { writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
resize: vi.fn().mockReturnThis(),
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { extractText } from "../../../packages/ai/src/ocr.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-image-data");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-ocr";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true, "text": "Hello World", "engine": "paddleocr"}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
text: "Hello World",
|
||||
engine: "paddleocr",
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
resize: vi.fn().mockReturnThis(),
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("extractText", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls ocr.py with input path and options JSON", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"ocr.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_ocr.png`, "{}"],
|
||||
expect.objectContaining({ timeout: expect.any(Number) }),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes quality option", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR, { quality: "best" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[1])).toEqual({ quality: "best" });
|
||||
});
|
||||
|
||||
it("serializes language option", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR, { language: "ja" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[1])).toEqual({ language: "ja" });
|
||||
});
|
||||
|
||||
it("serializes enhance option", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR, { enhance: true });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[1])).toEqual({ enhance: true });
|
||||
});
|
||||
|
||||
it("serializes deprecated engine option for backward compatibility", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR, { engine: "tesseract" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[1])).toEqual({ engine: "tesseract" });
|
||||
});
|
||||
|
||||
it("serializes all options together", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR, {
|
||||
quality: "fast",
|
||||
language: "en",
|
||||
enhance: false,
|
||||
});
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[1])).toEqual({
|
||||
quality: "fast",
|
||||
language: "en",
|
||||
enhance: false,
|
||||
});
|
||||
});
|
||||
|
||||
it("resizes input to max 2048px before writing", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
// Sharp is called with the input, then resize is called
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
});
|
||||
|
||||
it("writes resized PNG to outputDir", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
`${FAKE_OUTPUT_DIR}/input_ocr.png`,
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns OcrResult with text and engine", async () => {
|
||||
const result = await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
text: "Hello World",
|
||||
engine: "paddleocr",
|
||||
});
|
||||
});
|
||||
|
||||
it("returns text with special characters", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
text: "Price: $19.99\nDiscount: 15%",
|
||||
engine: "paddleocr",
|
||||
});
|
||||
|
||||
const result = await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.text).toBe("Price: $19.99\nDiscount: 15%");
|
||||
});
|
||||
|
||||
it("returns empty text string when no text detected", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
text: "",
|
||||
engine: "paddleocr",
|
||||
});
|
||||
|
||||
const result = await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.text).toBe("");
|
||||
});
|
||||
|
||||
it("returns engine information", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
text: "test",
|
||||
engine: "tesseract",
|
||||
});
|
||||
|
||||
const result = await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.engine).toBe("tesseract");
|
||||
});
|
||||
|
||||
it("returns undefined engine when not provided by Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
text: "test",
|
||||
});
|
||||
|
||||
const result = await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.engine).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("timeout calculation", () => {
|
||||
it("uses minimum 600000ms timeout for small images", async () => {
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
// 800x600 = 0.48 MP, 0.48 * 30 * 1000 = 14400 < 600000
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBe(600000);
|
||||
});
|
||||
|
||||
it("scales timeout for large images", async () => {
|
||||
// We need sharp to return large dimensions for the resized buffer
|
||||
// First call resizes the input, second call reads metadata of the resized buffer
|
||||
let callCount = 0;
|
||||
vi.mocked(sharp).mockImplementation(() => {
|
||||
callCount++;
|
||||
return {
|
||||
resize: vi.fn().mockReturnThis(),
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
metadata: vi.fn().mockResolvedValue({ width: 5000, height: 4000 }),
|
||||
} as unknown as ReturnType<typeof sharp>;
|
||||
});
|
||||
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
// 5000*4000 = 20 MP, 20 * 30 * 1000 = 600000 = 600000 (equal to min)
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.timeout).toBeGreaterThanOrEqual(600000);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "PaddleOCR initialization failed",
|
||||
});
|
||||
|
||||
await expect(extractText(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"PaddleOCR initialization failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(extractText(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("OCR failed");
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(extractText(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
|
||||
it("propagates parseStdoutJson errors", async () => {
|
||||
vi.mocked(parseStdoutJson).mockImplementation(() => {
|
||||
throw new Error("No JSON response from Python script");
|
||||
});
|
||||
|
||||
await expect(extractText(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"No JSON response from Python script",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await extractText(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"ocr.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,254 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { removeRedEye } from "../../../packages/ai/src/red-eye-removal.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-redeye-image");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-redeye";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 1,
|
||||
eyesCorrected: 2,
|
||||
width: 800,
|
||||
height: 600,
|
||||
format: "png",
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("removeRedEye", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls red_eye_removal.py with correct file paths", async () => {
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"red_eye_removal.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_redeye.png`, `${FAKE_OUTPUT_DIR}/output_redeye.png`, "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes sensitivity option", async () => {
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR, { sensitivity: 0.8 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ sensitivity: 0.8 });
|
||||
});
|
||||
|
||||
it("serializes strength option", async () => {
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR, { strength: 0.6 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ strength: 0.6 });
|
||||
});
|
||||
|
||||
it("serializes format and quality options", async () => {
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR, { format: "webp", quality: 90 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ format: "webp", quality: 90 });
|
||||
});
|
||||
|
||||
it("serializes all options together", async () => {
|
||||
const allOptions = {
|
||||
sensitivity: 0.9,
|
||||
strength: 0.7,
|
||||
format: "jpeg",
|
||||
quality: 85,
|
||||
};
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR, allOptions);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual(allOptions);
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns RedEyeRemovalResult with all fields", async () => {
|
||||
const result = await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
buffer: expect.any(Buffer),
|
||||
facesDetected: 1,
|
||||
eyesCorrected: 2,
|
||||
width: 800,
|
||||
height: 600,
|
||||
format: "png",
|
||||
});
|
||||
});
|
||||
|
||||
it("reads from default output path", async () => {
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_redeye.png`);
|
||||
});
|
||||
|
||||
it("reads from alternate output_path when provided", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
output_path: "/tmp/alt-redeye.webp",
|
||||
});
|
||||
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(readFile).toHaveBeenCalledWith("/tmp/alt-redeye.webp");
|
||||
});
|
||||
|
||||
it("defaults facesDetected to 0 when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.facesDetected).toBe(0);
|
||||
});
|
||||
|
||||
it("defaults eyesCorrected to 0 when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.eyesCorrected).toBe(0);
|
||||
});
|
||||
|
||||
it("defaults format to 'png' when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
});
|
||||
|
||||
const result = await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.format).toBe("png");
|
||||
});
|
||||
|
||||
it("reports zero corrections when no red eyes found", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 1,
|
||||
eyesCorrected: 0,
|
||||
width: 800,
|
||||
height: 600,
|
||||
format: "png",
|
||||
});
|
||||
|
||||
const result = await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.facesDetected).toBe(1);
|
||||
expect(result.eyesCorrected).toBe(0);
|
||||
});
|
||||
|
||||
it("handles multiple faces with multiple corrections", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
facesDetected: 3,
|
||||
eyesCorrected: 5,
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
format: "png",
|
||||
});
|
||||
|
||||
const result = await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.facesDetected).toBe(3);
|
||||
expect(result.eyesCorrected).toBe(5);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "Face detection model not available",
|
||||
});
|
||||
|
||||
await expect(removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Face detection model not available",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Red eye removal failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
|
||||
it("propagates bridge segfault", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process crashed (segmentation fault)"),
|
||||
);
|
||||
|
||||
await expect(removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("segmentation fault");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await removeRedEye(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"red_eye_removal.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,292 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { restorePhoto } from "../../../packages/ai/src/restoration.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-old-photo");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-restore";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
steps: ["denoise", "face_enhance"],
|
||||
scratchCoverage: 0.15,
|
||||
facesEnhanced: 2,
|
||||
isGrayscale: true,
|
||||
colorized: true,
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("restorePhoto", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls restore.py with correct file paths", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"restore.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_restore.png`, `${FAKE_OUTPUT_DIR}/output_restore.png`, "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes mode option", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, { mode: "heavy" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ mode: "heavy" });
|
||||
});
|
||||
|
||||
it("serializes scratchRemoval option", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, { scratchRemoval: true });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ scratchRemoval: true });
|
||||
});
|
||||
|
||||
it("serializes faceEnhancement option", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, { faceEnhancement: true });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ faceEnhancement: true });
|
||||
});
|
||||
|
||||
it("serializes fidelity option", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, { fidelity: 0.8 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ fidelity: 0.8 });
|
||||
});
|
||||
|
||||
it("serializes denoise and denoiseStrength options", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, { denoise: true, denoiseStrength: 0.5 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ denoise: true, denoiseStrength: 0.5 });
|
||||
});
|
||||
|
||||
it("serializes colorize option", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, { colorize: true });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ colorize: true });
|
||||
});
|
||||
|
||||
it("serializes all options together", async () => {
|
||||
const allOptions = {
|
||||
mode: "auto",
|
||||
scratchRemoval: true,
|
||||
faceEnhancement: true,
|
||||
fidelity: 0.8,
|
||||
denoise: true,
|
||||
denoiseStrength: 0.5,
|
||||
colorize: true,
|
||||
};
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, allOptions);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual(allOptions);
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns RestorePhotoResult with all fields populated", async () => {
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
buffer: expect.any(Buffer),
|
||||
width: 800,
|
||||
height: 600,
|
||||
steps: ["denoise", "face_enhance"],
|
||||
scratchCoverage: 0.15,
|
||||
facesEnhanced: 2,
|
||||
isGrayscale: true,
|
||||
colorized: true,
|
||||
});
|
||||
});
|
||||
|
||||
it("reads from default output path", async () => {
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_restore.png`);
|
||||
});
|
||||
|
||||
it("reads from alternate output_path when provided", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
output_path: "/tmp/alt-restore.webp",
|
||||
});
|
||||
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(readFile).toHaveBeenCalledWith("/tmp/alt-restore.webp");
|
||||
});
|
||||
|
||||
it("defaults steps to empty array when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 400,
|
||||
height: 300,
|
||||
});
|
||||
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.steps).toEqual([]);
|
||||
});
|
||||
|
||||
it("defaults scratchCoverage to 0 when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 400,
|
||||
height: 300,
|
||||
});
|
||||
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.scratchCoverage).toBe(0);
|
||||
});
|
||||
|
||||
it("defaults facesEnhanced to 0 when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 400,
|
||||
height: 300,
|
||||
});
|
||||
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.facesEnhanced).toBe(0);
|
||||
});
|
||||
|
||||
it("defaults isGrayscale to false when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 400,
|
||||
height: 300,
|
||||
});
|
||||
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.isGrayscale).toBe(false);
|
||||
});
|
||||
|
||||
it("defaults colorized to false when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 400,
|
||||
height: 300,
|
||||
});
|
||||
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.colorized).toBe(false);
|
||||
});
|
||||
|
||||
it("preserves multi-step restoration pipeline info", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 800,
|
||||
height: 600,
|
||||
steps: ["scratch_removal", "denoise", "face_enhance", "colorize"],
|
||||
scratchCoverage: 0.3,
|
||||
facesEnhanced: 4,
|
||||
isGrayscale: true,
|
||||
colorized: true,
|
||||
});
|
||||
|
||||
const result = await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.steps).toHaveLength(4);
|
||||
expect(result.scratchCoverage).toBe(0.3);
|
||||
expect(result.facesEnhanced).toBe(4);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "CodeFormer model weights not found",
|
||||
});
|
||||
|
||||
await expect(restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"CodeFormer model weights not found",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"Photo restoration failed",
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
|
||||
it("propagates OOM errors", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process killed (out of memory)"),
|
||||
);
|
||||
|
||||
await expect(restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("out of memory");
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await restorePhoto(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"restore.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,260 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("sharp", () => {
|
||||
const mockSharp = vi.fn(() => ({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}));
|
||||
return { default: mockSharp };
|
||||
});
|
||||
|
||||
vi.mock("node:fs/promises", () => ({
|
||||
readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("../../../packages/ai/src/bridge.js", () => ({
|
||||
runPythonWithProgress: vi.fn(),
|
||||
parseStdoutJson: vi.fn(),
|
||||
}));
|
||||
|
||||
import sharp from "sharp";
|
||||
import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
|
||||
import { upscale } from "../../../packages/ai/src/upscaling.js";
|
||||
|
||||
const FAKE_INPUT = Buffer.from("fake-small-image");
|
||||
const FAKE_OUTPUT_DIR = "/tmp/test-upscale";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
|
||||
vi.mocked(writeFile).mockResolvedValue(undefined);
|
||||
vi.mocked(runPythonWithProgress).mockResolvedValue({
|
||||
stdout: '{"success": true}',
|
||||
stderr: "",
|
||||
});
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 1600,
|
||||
height: 1200,
|
||||
method: "realesrgan",
|
||||
format: "png",
|
||||
});
|
||||
vi.mocked(sharp).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
png: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
|
||||
}) as unknown as ReturnType<typeof sharp>,
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("upscale", () => {
|
||||
describe("request serialization", () => {
|
||||
it("calls upscale.py with correct file paths", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"upscale.py",
|
||||
[`${FAKE_OUTPUT_DIR}/input_upscale.png`, `${FAKE_OUTPUT_DIR}/output_upscale.png`, "{}"],
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
|
||||
it("serializes scale option", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ scale: 4 });
|
||||
});
|
||||
|
||||
it("serializes model option", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { model: "realesrgan-x4plus-anime" });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ model: "realesrgan-x4plus-anime" });
|
||||
});
|
||||
|
||||
it("serializes faceEnhance option", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { faceEnhance: true });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ faceEnhance: true });
|
||||
});
|
||||
|
||||
it("serializes denoise option", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { denoise: 0.5 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ denoise: 0.5 });
|
||||
});
|
||||
|
||||
it("serializes format and quality options", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { format: "webp", quality: 90 });
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual({ format: "webp", quality: 90 });
|
||||
});
|
||||
|
||||
it("serializes all options together", async () => {
|
||||
const allOptions = {
|
||||
scale: 2,
|
||||
model: "realesrgan-x4plus",
|
||||
faceEnhance: true,
|
||||
denoise: 0.3,
|
||||
format: "jpeg",
|
||||
quality: 85,
|
||||
};
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, allOptions);
|
||||
|
||||
const args = vi.mocked(runPythonWithProgress).mock.calls[0][1];
|
||||
expect(JSON.parse(args[2])).toEqual(allOptions);
|
||||
});
|
||||
|
||||
it("converts input to PNG before writing", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(sharp).toHaveBeenCalledWith(FAKE_INPUT);
|
||||
expect(writeFile).toHaveBeenCalledWith(
|
||||
`${FAKE_OUTPUT_DIR}/input_upscale.png`,
|
||||
Buffer.from("mock-png-data"),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("response parsing", () => {
|
||||
it("returns UpscaleResult with all fields", async () => {
|
||||
const result = await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(result).toEqual({
|
||||
buffer: expect.any(Buffer),
|
||||
width: 1600,
|
||||
height: 1200,
|
||||
method: "realesrgan",
|
||||
format: "png",
|
||||
});
|
||||
});
|
||||
|
||||
it("reads from default output path", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
expect(readFile).toHaveBeenCalledWith(`${FAKE_OUTPUT_DIR}/output_upscale.png`);
|
||||
});
|
||||
|
||||
it("reads from alternate output_path when provided", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 1600,
|
||||
height: 1200,
|
||||
output_path: "/tmp/alt-upscale.webp",
|
||||
});
|
||||
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(readFile).toHaveBeenCalledWith("/tmp/alt-upscale.webp");
|
||||
});
|
||||
|
||||
it("defaults method to 'unknown' when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 1600,
|
||||
height: 1200,
|
||||
});
|
||||
|
||||
const result = await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.method).toBe("unknown");
|
||||
});
|
||||
|
||||
it("defaults format to 'png' when absent", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 1600,
|
||||
height: 1200,
|
||||
});
|
||||
|
||||
const result = await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
expect(result.format).toBe("png");
|
||||
});
|
||||
|
||||
it("returns correct dimensions for 4x upscale", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: true,
|
||||
width: 3200,
|
||||
height: 2400,
|
||||
method: "realesrgan",
|
||||
format: "png",
|
||||
});
|
||||
|
||||
const result = await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 });
|
||||
expect(result.width).toBe(3200);
|
||||
expect(result.height).toBe(2400);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("throws with custom error from Python", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({
|
||||
success: false,
|
||||
error: "RealESRGAN model file not found",
|
||||
});
|
||||
|
||||
await expect(upscale(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"RealESRGAN model file not found",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws fallback error when success: false without error string", async () => {
|
||||
vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
|
||||
|
||||
await expect(upscale(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("Upscaling failed");
|
||||
});
|
||||
|
||||
it("propagates bridge timeout", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(new Error("Python script timed out"));
|
||||
|
||||
await expect(upscale(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("timed out");
|
||||
});
|
||||
|
||||
it("propagates OOM errors", async () => {
|
||||
vi.mocked(runPythonWithProgress).mockRejectedValue(
|
||||
new Error("Process killed (out of memory)"),
|
||||
);
|
||||
|
||||
await expect(upscale(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("out of memory");
|
||||
});
|
||||
|
||||
it("propagates parseStdoutJson errors", async () => {
|
||||
vi.mocked(parseStdoutJson).mockImplementation(() => {
|
||||
throw new Error("No JSON response from Python script");
|
||||
});
|
||||
|
||||
await expect(upscale(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow(
|
||||
"No JSON response from Python script",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("onProgress forwarding", () => {
|
||||
it("passes onProgress to bridge", async () => {
|
||||
const onProgress = vi.fn();
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, {}, onProgress);
|
||||
|
||||
expect(runPythonWithProgress).toHaveBeenCalledWith(
|
||||
"upscale.py",
|
||||
expect.any(Array),
|
||||
expect.objectContaining({ onProgress }),
|
||||
);
|
||||
});
|
||||
|
||||
it("omits onProgress when not provided", async () => {
|
||||
await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR);
|
||||
|
||||
const options = vi.mocked(runPythonWithProgress).mock.calls[0][2];
|
||||
expect(options.onProgress).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user