pi-map/tests/llm-extract.test.ts

import { describe, it, expect } from "vitest";
import { extractFileLLM, extractFileHeuristic } from "../src/llm-extract.js";
import { writeFileSync, mkdtempSync } from "fs";
import { join } from "path";
import { tmpdir } from "os";
import type { LLMClient } from "../src/llm-client.js";

describe("llm-extract heuristics", () => {
	it("extracts TypeScript exports", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "test.ts");
		writeFileSync(
			file,
			`export function foo() {}
export class Bar {}
export const baz = 1;
export type Qux = string;
export { a, b as c };
`,
		);
		const result = await extractFileLLM(file);
		expect(result.exports).toContain("foo");
		expect(result.exports).toContain("Bar");
		expect(result.exports).toContain("baz");
		expect(result.exports).toContain("Qux");
		expect(result.exports).toContain("a");
		expect(result.exports).toContain("b");
	});

	it("extracts TypeScript imports", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "test.ts");
		writeFileSync(
			file,
			`import { foo } from "./bar";
import * as baz from "baz-lib";
import type { Qux } from "qux";
const x = require("legacy");
`,
		);
		const result = await extractFileLLM(file);
		expect(result.deps).toContain("./bar");
		expect(result.deps).toContain("baz-lib");
		expect(result.deps).toContain("qux");
		expect(result.deps).toContain("legacy");
	});

	it("infers purpose from filename patterns", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "userController.ts");
		writeFileSync(file, `export class UserController {}`);
		const result = await extractFileLLM(file);
		expect(result.purpose).toMatch(/Controller|Exports/);
	});

	it("handles non-code files", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "Dockerfile");
		writeFileSync(file, `FROM node:20\nWORKDIR /app`);
		const result = await extractFileLLM(file);
		expect(result.purpose).toBe("Container definition");
		expect(result.exports).toEqual([]);
	});
});

describe("llm-extract with mock client", () => {
	it("uses LLM client when provided", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "test.ts");
		writeFileSync(file, `export const foo = 1;`);

		const mockClient: LLMClient = {
			async complete() {
				return "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing";
			},
		};

		const result = await extractFileLLM(file, mockClient, tmpdir());
		expect(result.purpose).toBe("Test file");
		expect(result.deps).toEqual([]);
		expect(result.concepts).toContain("testing");
	});

	it("skips large files", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "big.ts");
		writeFileSync(file, "x".repeat(60 * 1024));

		const mockClient: LLMClient = {
			async complete() {
				return "PURPOSE: Should not call\nDEPS: none\nCONCEPTS: none";
			},
		};

		const result = await extractFileLLM(file, mockClient);
		expect(result.purpose).toBe("Large/generated file");
	});

	it("falls back to heuristics without client", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "utils.ts");
		writeFileSync(file, `export function helper() {}`);

		const result = await extractFileLLM(file);
		expect(result.purpose).toBe("Utility functions");
		expect(result.exports).toContain("helper");
	});
});

describe("extractFileHeuristic", () => {
	it("returns structured data", async () => {
		const dir = mkdtempSync(join(tmpdir(), "pi-map-"));
		const file = join(dir, "test.ts");
		writeFileSync(file, `export const x = 1;`);

		const result = await extractFileHeuristic(file);
		expect(result.purpose).toBeDefined();
		expect(Array.isArray(result.exports)).toBe(true);
		expect(Array.isArray(result.deps)).toBe(true);
		expect(Array.isArray(result.concepts)).toBe(true);
	});
});