import type { Rect } from "@xyflow/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { exportGraphAsJSON, exportGraphAsPNG } from "@/lib/attack-paths/export"; import type { AttackPathGraphData } from "@/types/attack-paths"; const bounds: Rect = { x: 0, y: 0, width: 100, height: 100 }; const graphData: AttackPathGraphData = { nodes: [ { id: "internet", labels: ["Internet"], properties: { name: "Internet" } }, { id: "ec2-1", labels: ["EC2Instance"], properties: { name: "api-server-01" }, }, ], edges: [ { id: "edge-1", source: "internet", target: "ec2-1", type: "CAN_REACH" }, ], }; const buildContainerWithViewport = () => { const container = document.createElement("div"); const reactFlow = document.createElement("div"); reactFlow.className = "react-flow"; const viewport = document.createElement("div"); viewport.className = "react-flow__viewport"; reactFlow.appendChild(viewport); container.appendChild(reactFlow); return container; }; class TestImage { onload: (() => void) | null = null; onerror: (() => void) | null = null; decoding = "async"; set src(_value: string) { queueMicrotask(() => this.onload?.()); } } describe("exportGraphAsPNG", () => { beforeEach(() => { vi.restoreAllMocks(); vi.stubGlobal("Image", TestImage); vi.spyOn(URL, "createObjectURL").mockReturnValue("blob:graph-export"); vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {}); vi.spyOn(HTMLCanvasElement.prototype, "getContext").mockReturnValue({ arc: vi.fn(), beginPath: vi.fn(), bezierCurveTo: vi.fn(), closePath: vi.fn(), fill: vi.fn(), fillRect: vi.fn(), fillText: vi.fn(), lineTo: vi.fn(), moveTo: vi.fn(), quadraticCurveTo: vi.fn(), restore: vi.fn(), save: vi.fn(), setLineDash: vi.fn(), stroke: vi.fn(), } as unknown as CanvasRenderingContext2D); vi.spyOn(HTMLCanvasElement.prototype, "toDataURL").mockReturnValue( "data:image/png;base64,AAAA", ); }); it("throws when the container is not mounted", async () => { await expect( exportGraphAsPNG(null, bounds, undefined, graphData), ).rejects.toThrow("Graph container not mounted"); }); it("throws when the React Flow root is missing inside the container", async () => { const container = document.createElement("div"); await expect( exportGraphAsPNG(container, bounds, undefined, graphData), ).rejects.toThrow("React Flow root not found in container"); }); it("throws when the React Flow viewport is missing inside the container", async () => { const container = document.createElement("div"); const reactFlow = document.createElement("div"); reactFlow.className = "react-flow"; container.appendChild(reactFlow); await expect( exportGraphAsPNG(container, bounds, undefined, graphData), ).rejects.toThrow("React Flow viewport not found in container"); }); it("throws when bounds are null (no nodes to export)", async () => { const container = buildContainerWithViewport(); await expect( exportGraphAsPNG(container, null, undefined, graphData), ).rejects.toThrow("No nodes to export"); }); it("throws when graph data has no nodes", async () => { const container = buildContainerWithViewport(); await expect( exportGraphAsPNG(container, bounds, undefined, { nodes: [] }), ).rejects.toThrow("No nodes to export"); }); it("renders graph data to a PNG download", async () => { const container = buildContainerWithViewport(); const appendChild = vi.spyOn(document.body, "appendChild"); await exportGraphAsPNG(container, bounds, "graph.png", graphData); expect(HTMLCanvasElement.prototype.toDataURL).toHaveBeenCalledWith( "image/png", ); const link = appendChild.mock.calls.find( ([element]) => element instanceof HTMLAnchorElement, )?.[0] as HTMLAnchorElement | undefined; expect(link?.download).toBe("graph.png"); expect(link?.href).toBe("data:image/png;base64,AAAA"); }); it("renders exported long resource labels with the same wrapping as graph nodes", async () => { const container = buildContainerWithViewport(); const longLabelGraphData: AttackPathGraphData = { nodes: [ { id: "role-1", labels: ["AWSRole"], properties: { name: "AWSReservedSSO_AdministratorAccessExtra" }, }, ], }; await exportGraphAsPNG(container, bounds, "graph.png", longLabelGraphData); const context = vi.mocked(HTMLCanvasElement.prototype.getContext).mock .results[0]?.value as CanvasRenderingContext2D; const fillText = vi.mocked(context.fillText); expect(fillText).toHaveBeenCalledWith( "AWSReservedSSO_A", expect.any(Number), expect.any(Number), 136, ); expect(fillText).toHaveBeenCalledWith( "dministratorAcce", expect.any(Number), expect.any(Number), 136, ); expect(fillText).toHaveBeenCalledWith( "ssExtra", expect.any(Number), expect.any(Number), 136, ); expect(fillText).toHaveBeenCalledWith( "AWS Role", expect.any(Number), expect.any(Number), 136, ); }); it("re-throws a generic export error when canvas is unavailable", async () => { const container = buildContainerWithViewport(); vi.spyOn(HTMLCanvasElement.prototype, "getContext").mockReturnValue(null); const consoleError = vi .spyOn(console, "error") .mockImplementation(() => {}); await expect( exportGraphAsPNG(container, bounds, undefined, graphData), ).rejects.toThrow("Failed to export graph"); expect(consoleError).toHaveBeenCalled(); }); }); describe("exportGraphAsJSON", () => { it("re-throws a generic export error when serialization fails", () => { const circular: Record = {}; circular.self = circular; const consoleError = vi .spyOn(console, "error") .mockImplementation(() => {}); expect(() => exportGraphAsJSON(circular)).toThrow("Failed to export graph"); expect(consoleError).toHaveBeenCalled(); consoleError.mockRestore(); }); });