import { chmodSync, copyFileSync, mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync, } from "node:fs"; import { tmpdir } from "node:os"; import { delimiter, join } from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { createTestSession, type TestSession } from "../../../test/harness"; import uvExtension, { __test } from "../index"; let t: TestSession | undefined; let fakeBin: string | undefined; afterEach(() => { t?.dispose(); if (fakeBin) rmSync(fakeBin, { recursive: true, force: true }); vi.unstubAllEnvs(); t = undefined; fakeBin = undefined; }); describe("uv extension", () => { it("loads in real Pi extension runtime", async () => { t = await createTestSession({ extensionFactories: [uvExtension] }); expect(t.session.extensionRunner.extensions).toHaveLength(1); const bash = t.session.extensionRunner.extensions[0].tools.get("bash").definition; expect(bash.description).toContain("routes PATH-based Python through uv"); expect(bash.promptSnippet).toContain("Execute bash commands"); expect(bash.promptGuidelines).toEqual([ expect.stringContaining("Use bash to inspect PI_*"), expect.stringContaining("Use bash with uv instead of pip"), ]); }); it("records availability outcome in its debug sandbox when enabled", async () => { fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-empty-bin-")); vi.stubEnv("PATH", fakeBin); const extensionPath = join(import.meta.dirname, "../index.ts"); t = await createTestSession({ env: { PI_UV_DEBUG: "1" }, extensions: [extensionPath], }); const extension = t.session.extensionRunner.extensions[0]; const context = t.session.extensionRunner.createCommandContext(); for (const handler of extension.handlers.get("session_start")) await handler({ type: "session_start" }, context); const directory = join(t.cwd, ".test-state", "pi-ext", "debug", "uv"); await vi.waitFor(() => { const lines = readdirSync(directory).flatMap((file) => readFileSync(join(directory, file), "utf8").trim().split("\n"), ); expect( lines.some((line) => line.includes("availability.check.finish")), ).toBe(true); }); for (const handler of extension.handlers.get("session_shutdown")) await handler({ type: "session_shutdown" }, context); const records = readFileSync( join(directory, readdirSync(directory)[0]), "utf8", ) .trim() .split("\n") .map((line) => JSON.parse(line)); expect(records.map((record) => record.event)).toEqual( expect.arrayContaining([ "session.start", "session.shutdown", "availability.check.finish", ]), ); expect(records).toContainEqual( expect.objectContaining({ event: "availability.check.finish", installed: false, }), ); }); it("detects a controlled executable on PATH through the extension", async () => { fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-bin-")); if (process.platform === "win32") { copyFileSync(process.execPath, join(fakeBin, "uv.exe")); } else { const uv = join(fakeBin, "uv"); writeFileSync(uv, "#!/bin/sh\nexit 0\n", "utf8"); chmodSync(uv, 0o755); } expect(await __test.uvInstalled({ PATH: fakeBin }, process.platform)).toBe( true, ); t = await createTestSession({ env: { PATH: fakeBin }, extensionFactories: [uvExtension], }); expect(t.events.uiCallsFor("notify")).toEqual([]); }); it("warns through the extension when PATH has no uv executable", async () => { fakeBin = mkdtempSync(join(tmpdir(), "pi-uv-bin-")); t = await createTestSession({ env: { PATH: fakeBin }, extensionFactories: [uvExtension], }); await vi.waitFor(() => expect(t?.events.uiCallsFor("notify").at(-1)?.args).toEqual([ "uv extension: uv not found on PATH", "warning", ]), ); }); it("blocks package-manager commands with uv guidance", () => { expect(__test.getBlockedCommandMessage("pip install rich")).toContain( "uv add PACKAGE", ); expect(__test.getBlockedCommandMessage("pip3 install rich")).toContain( "pip3 is disabled", ); expect(__test.getBlockedCommandMessage("poetry add rich")).toContain( "uv sync", ); expect( __test.getBlockedCommandMessage("python -m pip install rich"), ).toContain("python -m pip"); }); it("blocks venv and py_compile bytecode writes", () => { expect(__test.getBlockedCommandMessage("python3 -m venv .venv")).toContain( "uv venv", ); expect( __test.getBlockedCommandMessage(".venv/bin/python -m py_compile x.py"), ).toContain("__pycache__"); }); it("leaves normal Python execution alone", () => { expect(__test.getBlockedCommandMessage("python script.py")).toBeNull(); expect( __test.getBlockedCommandMessage("uv run python script.py"), ).toBeNull(); }); it("prepends shell shims to Unix PATH without mutating input", () => { const ctx = { command: "python --version", cwd: "/repo", env: { PATH: "/usr/bin", TOKEN: "preserved" }, }; const rewritten = __test.prependInterceptedCommands(ctx, "linux"); expect(rewritten).not.toBe(ctx); expect(rewritten.env.PATH).toBe( `${join(import.meta.dirname, "../intercepted-commands")}${delimiter}/usr/bin`, ); expect(rewritten.env.TOKEN).toBe("preserved"); expect(ctx.env.PATH).toBe("/usr/bin"); }); it("does not expose extensionless shell shims to Windows", () => { const ctx = { command: "python --version", cwd: "C:/repo", env: { Path: "C:/Windows" }, }; expect(__test.prependInterceptedCommands(ctx, "win32")).toBe(ctx); }); });