238 lines
9.2 KiB
TypeScript
238 lines
9.2 KiB
TypeScript
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(
|
|
'<svg xmlns="http://www.w3.org/2000/svg"><rect width="100" height="50"/></svg>',
|
|
);
|
|
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<string, unknown>) => Promise<import("../services/renderers/types.js").RenderResult>;
|
|
|
|
beforeEach(async () => {
|
|
vi.clearAllMocks();
|
|
// Restore default mocks
|
|
mockPuterChatComplete.mockResolvedValue(MOCK_MERMAID_SOURCE);
|
|
mockChromiumLaunch.mockResolvedValue(mockBrowser);
|
|
mockBrowser.newPage.mockResolvedValue(mockPage);
|
|
mockPageEval.mockResolvedValue(
|
|
'<svg xmlns="http://www.w3.org/2000/svg"><rect width="100" height="50"/></svg>',
|
|
);
|
|
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);
|
|
});
|
|
});
|