import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import type { DiagramBundle } from "../services/renderers/types.js"; // ─── Mock playwright-core chromium ───────────────────────────────────────────── const mockPageSetContent = vi.fn().mockResolvedValue(undefined); const mockPageWaitForSelector = vi.fn().mockResolvedValue(undefined); const mockPageEval = vi.fn().mockResolvedValue( '', ); const mockBrowserClose = vi.fn().mockResolvedValue(undefined); const mockPage = { setContent: mockPageSetContent, waitForSelector: mockPageWaitForSelector, $eval: mockPageEval, }; const mockBrowser = { newPage: vi.fn().mockResolvedValue(mockPage), close: mockBrowserClose, }; const mockChromiumLaunch = vi.fn().mockResolvedValue(mockBrowser); vi.mock("playwright-core", () => ({ chromium: { launch: mockChromiumLaunch, }, })); // ─── Mock LLM inference ───────────────────────────────────────────────────────── const MOCK_MERMAID_SOURCE = "graph TD\n A[Login]-->B[Validate]\n B-->C[Dashboard]"; const mockPuterChatComplete = vi.fn().mockResolvedValue(MOCK_MERMAID_SOURCE); vi.mock("../services/puter-inference.js", () => ({ puterChatComplete: (...args: unknown[]) => mockPuterChatComplete(...args), })); // ─── Tests ────────────────────────────────────────────────────────────────────── describe("stripUnsafeDirectives", () => { let stripUnsafeDirectives: (source: string) => { cleaned: string; stripped: boolean }; beforeEach(async () => { // Re-import to get a fresh module each test const mod = await import("../services/renderers/diagram-renderer.js"); stripUnsafeDirectives = mod.stripUnsafeDirectives; }); it("strips %%{init}%% blocks and marks stripped=true", () => { const input = '%%{init: {"theme": "dark"}}%%\ngraph TD\n A-->B'; const result = stripUnsafeDirectives(input); expect(result.stripped).toBe(true); expect(result.cleaned).not.toContain("%%{"); expect(result.cleaned).toContain("graph TD"); }); it('strips click "url" lines and marks stripped=true', () => { const input = 'graph TD\n A-->B\n click A "https://evil.com"'; const result = stripUnsafeDirectives(input); expect(result.stripped).toBe(true); expect(result.cleaned).not.toContain("click A"); expect(result.cleaned).toContain("graph TD"); }); it("strips click call fn() lines and marks stripped=true", () => { const input = "graph TD\n A-->B\n click A call myFn()"; const result = stripUnsafeDirectives(input); expect(result.stripped).toBe(true); expect(result.cleaned).not.toContain("click A"); }); it("leaves clean source unchanged with stripped=false", () => { const input = "graph TD\n A-->B"; const result = stripUnsafeDirectives(input); expect(result.stripped).toBe(false); expect(result.cleaned).toBe("graph TD\n A-->B"); }); it("strips both init and click directives simultaneously", () => { const input = '%%{init: {"theme": "dark"}}%%\ngraph TD\n A-->B\n click A "https://evil.com"'; const result = stripUnsafeDirectives(input); expect(result.stripped).toBe(true); expect(result.cleaned).not.toContain("%%{"); expect(result.cleaned).not.toContain("click A"); expect(result.cleaned).toContain("graph TD"); }); }); describe("buildDiagramPrompt", () => { let buildDiagramPrompt: ( description: string, diagramType: string, ) => { system: string; user: string }; beforeEach(async () => { const mod = await import("../services/renderers/diagram-renderer.js"); buildDiagramPrompt = mod.buildDiagramPrompt; }); it("includes 'flowchart' when diagramType is flowchart", () => { const result = buildDiagramPrompt("a login flow", "flowchart"); const combined = result.system + result.user; expect(combined.toLowerCase()).toContain("flowchart"); }); it("includes 'architecture' when diagramType is architecture", () => { const result = buildDiagramPrompt("a microservices setup", "architecture"); const combined = result.system + result.user; expect(combined.toLowerCase()).toContain("architecture"); }); it("includes the user's natural language description in the user prompt", () => { const desc = "unique description for test xyz123"; const result = buildDiagramPrompt(desc, "flowchart"); expect(result.user).toContain(desc); }); it("instructs LLM to output ONLY valid Mermaid syntax without markdown fences", () => { const result = buildDiagramPrompt("some flow", "flowchart"); const system = result.system.toLowerCase(); expect(system).toContain("mermaid"); expect(system).toMatch(/only|no markdown|no explanation/i); }); it("includes sequence-specific preamble for sequence diagrams", () => { const result = buildDiagramPrompt("an API request flow", "sequence"); const combined = result.system + result.user; expect(combined.toLowerCase()).toContain("sequence"); }); it("includes ERD-specific preamble for erd diagrams", () => { const result = buildDiagramPrompt("user-post relationship", "erd"); const combined = result.system + result.user; expect(combined.toLowerCase()).toContain("erd"); }); it("includes mindmap-specific preamble for mindmap diagrams", () => { const result = buildDiagramPrompt("project ideas", "mindmap"); const combined = result.system + result.user; expect(combined.toLowerCase()).toContain("mindmap"); }); }); describe("renderDiagram integration", () => { let renderDiagram: (input: Record) => Promise; beforeEach(async () => { vi.clearAllMocks(); // Restore default mocks mockPuterChatComplete.mockResolvedValue(MOCK_MERMAID_SOURCE); mockChromiumLaunch.mockResolvedValue(mockBrowser); mockBrowser.newPage.mockResolvedValue(mockPage); mockPageEval.mockResolvedValue( '', ); mockBrowserClose.mockResolvedValue(undefined); const mod = await import("../services/renderers/diagram-renderer.js"); renderDiagram = mod.renderDiagram; }); afterEach(() => { vi.clearAllMocks(); }); it("returns a RenderResult with diagram-bundle JSON structure", async () => { const result = await renderDiagram({ prompt: "A login flow with validation", diagramType: "flowchart", darkMode: false, }); expect(result.filename).toBe("diagram-bundle.json"); expect(result.contentType).toBe("application/json"); const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle; expect(bundle.type).toBe("diagram-bundle"); expect(typeof bundle.svgBase64).toBe("string"); expect(typeof bundle.pngBase64).toBe("string"); expect(typeof bundle.mermaidSource).toBe("string"); expect(typeof bundle.stripped).toBe("boolean"); }); it("mermaidSource in bundle is the LLM-generated Mermaid, not the original prompt", async () => { const result = await renderDiagram({ prompt: "A login flow with validation", diagramType: "flowchart", darkMode: false, }); const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle; // The mermaidSource should contain the LLM-returned Mermaid, not the English prompt expect(bundle.mermaidSource).toContain("graph TD"); expect(bundle.mermaidSource).not.toBe("A login flow with validation"); }); it("calls browser.close() for cleanup", async () => { await renderDiagram({ prompt: "test prompt", diagramType: "flowchart", darkMode: false, }); expect(mockBrowserClose).toHaveBeenCalled(); }); it("calls browser.close() even when rendering throws", async () => { mockPageWaitForSelector.mockRejectedValueOnce(new Error("timeout")); await expect( renderDiagram({ prompt: "bad diagram", diagramType: "flowchart", darkMode: false }), ).rejects.toThrow(); expect(mockBrowserClose).toHaveBeenCalled(); }); it("uses source directly when provided (re-render path, skips LLM)", async () => { mockPuterChatComplete.mockClear(); const result = await renderDiagram({ source: "graph TD\n X-->Y", diagramType: "flowchart", darkMode: false, }); // LLM should NOT have been called expect(mockPuterChatComplete).not.toHaveBeenCalled(); const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle; expect(bundle.mermaidSource).toContain("graph TD"); }); it("sets stripped=true in bundle when LLM-generated source contains click directives", async () => { mockPuterChatComplete.mockResolvedValueOnce( 'graph TD\n A-->B\n click A "https://evil.com"', ); const result = await renderDiagram({ prompt: "diagram with click", diagramType: "flowchart", darkMode: false, }); const bundle = JSON.parse(result.buffer.toString()) as DiagramBundle; expect(bundle.stripped).toBe(true); }); });