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")), 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), })); 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")), metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), }) as unknown as ReturnType, ); }); 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"); }); it("propagates OOM errors from bridge", async () => { vi.mocked(runPythonWithProgress).mockRejectedValue( new Error("Process killed (out of memory)"), ); await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("out of memory"); }); }); 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 }), ); }); it("omits onProgress when not provided", async () => { await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR); const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; expect(options.onProgress).toBeUndefined(); }); }); describe("timeout calculation", () => { it("uses minimum 300000ms timeout for small images", async () => { await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR); // 800x600 = 0.48 MP, 0.48 * 120000 = 57600 < 300000 const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; expect(options.timeout).toBe(300000); }); it("scales timeout for large images using megapixels * 120000", async () => { vi.mocked(sharp).mockImplementation( () => ({ png: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), metadata: vi.fn().mockResolvedValue({ width: 4000, height: 3000 }), }) as unknown as ReturnType, ); await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR); // 4000x3000 = 12 MP, 12 * 120000 = 1440000 > 300000 const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; expect(options.timeout).toBe(1440000); }); it("uses metadata from the PNG-converted buffer for timeout", async () => { // The second sharp() call reads metadata from the PNG buffer let callCount = 0; vi.mocked(sharp).mockImplementation( () => ({ png: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), metadata: vi .fn() .mockResolvedValue( callCount++ === 0 ? { width: 800, height: 600 } : { width: 2000, height: 2000 }, ), }) as unknown as ReturnType, ); await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR); // Timeout is based on the metadata call, which returns dimensions const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; expect(options.timeout).toBeGreaterThanOrEqual(300000); }); }); describe("parseStdoutJson error propagation", () => { it("propagates parseStdoutJson errors", async () => { vi.mocked(parseStdoutJson).mockImplementation(() => { throw new Error("No JSON response from Python script"); }); await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow( "No JSON response from Python script", ); }); }); describe("sharp conversion", () => { it("converts input to PNG and writes to outputDir", async () => { await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR); expect(sharp).toHaveBeenCalledWith(FAKE_INPUT); expect(writeFile).toHaveBeenCalledWith( `${FAKE_OUTPUT_DIR}/input_denoise.png`, Buffer.from("mock-png-data"), ); }); it("propagates sharp toBuffer error", async () => { vi.mocked(sharp).mockImplementation( () => ({ png: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockRejectedValue(new Error("Image decode failed")), metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), }) as unknown as ReturnType, ); await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow( "Image decode failed", ); }); }); describe("edge cases", () => { it("handles zero dimensions from metadata", async () => { vi.mocked(sharp).mockImplementation( () => ({ png: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), metadata: vi.fn().mockResolvedValue({ width: undefined, height: undefined }), }) as unknown as ReturnType, ); await noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR); // Should not throw; dimensions default to 0 const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; expect(options.timeout).toBe(300000); // min timeout since megapixels = 0 }); it("propagates segfault from bridge", async () => { vi.mocked(runPythonWithProgress).mockRejectedValue( new Error("Process crashed (segmentation fault)"), ); await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("segmentation fault"); }); it("propagates readFile error after successful Python run", async () => { vi.mocked(readFile).mockRejectedValueOnce(new Error("ENOENT: output missing")); await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("output missing"); }); it("propagates writeFile error", async () => { vi.mocked(writeFile).mockRejectedValueOnce(new Error("ENOSPC: disk full")); await expect(noiseRemoval(FAKE_INPUT, FAKE_OUTPUT_DIR)).rejects.toThrow("disk full"); }); }); });