pi-map/tests/pi-extension.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";

// Mock child_process before importing the extension
vi.mock("child_process", () => ({
	execSync: vi.fn(),
}));

vi.mock("fs", async () => {
	const actual = await vi.importActual<typeof import("fs")>("fs");
	return {
		...actual,
		readFileSync: vi.fn(),
	};
});

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 }),
	},
}));

import extension from "../pi-extension.js";
import { execSync } from "child_process";
import { readFileSync } from "fs";

describe("pi-extension", () => {
	let registeredTools: Record<string, any> = {};
	let registeredEvents: Record<string, any> = {};
	let mockCtx: any;
	let mockNotify: ReturnType<typeof vi.fn>;

	beforeEach(() => {
		registeredTools = {};
		registeredEvents = {};
		mockNotify = vi.fn();
		mockCtx = {
			cwd: "/home/project",
			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);
		vi.clearAllMocks();
	});

	afterEach(() => {
		vi.restoreAllMocks();
	});

	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("calls npx project-map init with default cwd", async () => {
			vi.mocked(execSync).mockReturnValue("Generated 5 .pi-map.md files\n");
			const tool = registeredTools.project_map_init;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(execSync).toHaveBeenCalledWith(
				expect.stringContaining("npx project-map init"),
				expect.objectContaining({ cwd: "/home/project" }),
			);
			expect(result.details.success).toBe(true);
		});

		it("calls npx project-map init with custom path", async () => {
			vi.mocked(execSync).mockReturnValue("Generated 3 .pi-map.md files\n");
			const tool = registeredTools.project_map_init;
			const result = await tool.execute(
				"tool-1",
				{ path: "./src" },
				null,
				null,
				mockCtx,
			);
			expect(execSync).toHaveBeenCalledWith(
				expect.stringContaining("npx project-map init ./src"),
				expect.anything(),
			);
			expect(result.details.success).toBe(true);
		});

		it("reports failure on error", async () => {
			vi.mocked(execSync).mockImplementation(() => {
				const err = new Error("Command failed") as any;
				err.stdout = "";
				err.stderr = "No such file";
				throw err;
			});
			const tool = registeredTools.project_map_init;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.success).toBe(false);
		});
	});

	describe("project_map_patch tool", () => {
		it("calls npx project-map patch with file path", async () => {
			vi.mocked(execSync).mockReturnValue("Patched src/foo.ts\n");
			const tool = registeredTools.project_map_patch;
			const result = await tool.execute(
				"tool-1",
				{ file_path: "src/components/Button.tsx" },
				null,
				null,
				mockCtx,
			);
			expect(execSync).toHaveBeenCalledWith(
				expect.stringContaining(
					"npx project-map patch src/components/Button.tsx",
				),
				expect.anything(),
			);
			expect(result.details.success).toBe(true);
		});
	});

	describe("project_map_validate tool", () => {
		it("calls npx project-map validate and reports clean", async () => {
			vi.mocked(execSync).mockReturnValue("All .pi-map.md files are clean.\n");
			const tool = registeredTools.project_map_validate;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(execSync).toHaveBeenCalledWith(
				expect.stringContaining("npx project-map validate"),
				expect.anything(),
			);
			expect(result.details.clean).toBe(true);
		});

		it("reports not clean when output lacks 'clean'", async () => {
			vi.mocked(execSync).mockReturnValue("Found 2 discrepancies\n");
			const tool = registeredTools.project_map_validate;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.clean).toBe(false);
		});
	});

	describe("project_map_reinit tool", () => {
		it("calls npx project-map reinit with path", async () => {
			vi.mocked(execSync).mockReturnValue("Regenerated 7 files\n");
			const tool = registeredTools.project_map_reinit;
			const result = await tool.execute(
				"tool-1",
				{ path: "./src" },
				null,
				null,
				mockCtx,
			);
			expect(execSync).toHaveBeenCalledWith(
				expect.stringContaining("npx project-map reinit ./src"),
				expect.anything(),
			);
			expect(result.details.success).toBe(true);
		});
	});

	describe("session_start event", () => {
		it("notifies when dirty .pi-map.md files exist", async () => {
			vi.mocked(execSync).mockReturnValueOnce(
				"./src/.pi-map.md\n./tests/.pi-map.md\n",
			); // findPiMapFiles
			vi.mocked(readFileSync)
				.mockReturnValueOnce("## dirty\n2024-01-01: patched\n") // dirty
				.mockReturnValueOnce("## dirty\n-\n"); // clean

			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 () => {
			vi.mocked(execSync).mockReturnValueOnce("./src/.pi-map.md\n");
			vi.mocked(readFileSync).mockReturnValueOnce("## dirty\n-\n");

			const handler = registeredEvents.session_start;
			await handler(null, mockCtx);

			expect(mockNotify).not.toHaveBeenCalled();
		});

		it("does nothing when no maps exist", async () => {
			vi.mocked(execSync).mockImplementation(() => {
				throw new Error("no maps");
			});

			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 () => {
			vi.mocked(execSync).mockReturnValueOnce("./src/.pi-map.md\n");
			vi.mocked(readFileSync).mockReturnValueOnce("content");

			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 () => {
			vi.mocked(execSync).mockImplementation(() => {
				throw new Error("no maps");
			});

			const handler = registeredEvents.before_agent_start;
			const result = await handler(null, mockCtx);

			expect(result).toEqual({});
		});
	});
});