import { describe, it, expect, vi, beforeEach } from "vitest"; import { mkdtempSync, writeFileSync, existsSync, unlinkSync } from "fs"; import { join } from "path"; import { tmpdir, homedir } from "os"; vi.mock("@mariozechner/pi-coding-agent", () => ({ ExtensionAPI: class {}, })); vi.mock("typebox", () => ({ Type: { Object: (props: unknown) => props, Optional: (prop: unknown) => prop, String: (opts: unknown) => ({ type: "string", ...(opts as object) }), }, })); // Mock Pi's AI module so PiLLMClient works in tests vi.mock("@mariozechner/pi-ai", () => ({ complete: vi.fn(async (_model: any, _context: any) => ({ role: "assistant" as const, content: [ { type: "text", text: "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing", }, ], })), })); import extension from "../pi-extension.js"; const CACHE_FILE = join( homedir(), ".cache", "pi-project-map", "llm-cache.json", ); function clearCache() { if (existsSync(CACHE_FILE)) unlinkSync(CACHE_FILE); } describe("pi-extension", () => { let registeredTools: Record = {}; let registeredEvents: Record = {}; let mockCtx: any; let mockNotify: ReturnType; beforeEach(() => { clearCache(); registeredTools = {}; registeredEvents = {}; mockNotify = vi.fn(); mockCtx = { cwd: "/home/project", model: { provider: "openai", id: "gpt-4o-mini", api: "openai-completions", }, modelRegistry: { getApiKeyAndHeaders: vi.fn(async () => ({ apiKey: "test-key", headers: {}, })), }, ui: { notify: mockNotify }, }; const mockPi = { registerTool: vi.fn((tool: any) => { registeredTools[tool.name] = tool; }), on: vi.fn((event: string, handler: any) => { registeredEvents[event] = handler; }), }; extension(mockPi); }); describe("tool registration", () => { it("registers 4 tools", () => { expect(Object.keys(registeredTools)).toHaveLength(4); expect(registeredTools).toHaveProperty("project_map_init"); expect(registeredTools).toHaveProperty("project_map_patch"); expect(registeredTools).toHaveProperty("project_map_validate"); expect(registeredTools).toHaveProperty("project_map_reinit"); }); it("registers session_start and before_agent_start events", () => { expect(registeredEvents).toHaveProperty("session_start"); expect(registeredEvents).toHaveProperty("before_agent_start"); }); }); describe("project_map_init tool", () => { it("generates maps with mocked Pi LLM", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`); mockCtx.cwd = dir; const tool = registeredTools.project_map_init; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.success).toBe(true); expect(result.content[0].text).toContain("Generated"); }); it("returns error when no Pi model configured", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`); mockCtx.cwd = dir; mockCtx.model = null; const tool = registeredTools.project_map_init; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.success).toBe(false); expect(result.content[0].text).toContain("no model configured"); }); }); describe("project_map_patch tool", () => { it("patches map with mocked Pi LLM", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); const file = join(dir, "test.ts"); writeFileSync(file, `export const x = ${Date.now()};`); mockCtx.cwd = dir; const tool = registeredTools.project_map_patch; const result = await tool.execute( "tool-1", { file_path: file }, null, null, mockCtx, ); expect(result.details.success).toBe(true); expect(result.content[0].text).toContain("Patched map"); }); }); describe("project_map_validate tool", () => { it("reports clean when map exists and is up to date", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync( join(dir, ".pi-map.md"), "# .\n## role\nTest\n## files\n## arch\n## dirty\n-\n", ); mockCtx.cwd = dir; const tool = registeredTools.project_map_validate; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.clean).toBe(true); expect(result.content[0].text).toContain("clean"); }); it("reports missing when files exist but no map", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, "test.ts"), "export const x = 1;"); mockCtx.cwd = dir; const tool = registeredTools.project_map_validate; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.clean).toBe(false); expect(result.content[0].text).toContain("missing"); }); }); describe("project_map_reinit tool", () => { it("regenerates maps with mocked Pi LLM", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`); mockCtx.cwd = dir; const tool = registeredTools.project_map_reinit; const result = await tool.execute("tool-1", {}, null, null, mockCtx); expect(result.details.success).toBe(true); expect(result.content[0].text).toContain("Regenerated"); }); }); describe("session_start event", () => { it("notifies when dirty .pi-map.md files exist", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync( join(dir, ".pi-map.md"), "# .\n## dirty\n2024-01-01: patched\n", ); mockCtx.cwd = dir; const handler = registeredEvents.session_start; await handler(null, mockCtx); expect(mockNotify).toHaveBeenCalledWith( expect.stringContaining("1 dirty packages detected"), "warning", ); }); it("does not notify when all clean", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, ".pi-map.md"), "# .\n## dirty\n-\n"); mockCtx.cwd = dir; const handler = registeredEvents.session_start; await handler(null, mockCtx); expect(mockNotify).not.toHaveBeenCalled(); }); it("does nothing when no maps exist", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); mockCtx.cwd = dir; const handler = registeredEvents.session_start; await handler(null, mockCtx); expect(mockNotify).not.toHaveBeenCalled(); }); }); describe("before_agent_start event", () => { it("injects hint when .pi-map.md files exist", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n"); mockCtx.cwd = dir; const handler = registeredEvents.before_agent_start; const result = await handler(null, mockCtx); expect(result).toHaveProperty("message"); expect(result.message.content).toContain("project_map_patch"); expect(result.message.display).toBe(false); }); it("returns empty object when no maps exist", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); mockCtx.cwd = dir; const handler = registeredEvents.before_agent_start; const result = await handler(null, mockCtx); expect(result).toEqual({}); }); }); });