my-pi/pi-worktree/test/state.test.ts

import assert from "node:assert/strict";
import { existsSync } from "node:fs";
import { mkdir, mkdtemp, realpath, stat, symlink, unlink, utimes, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join } from "node:path";
import test from "node:test";

import {
	addEntry,
	findByCwd,
	findBySlug,
	loadState,
	removeEntry,
	saveState,
	stateFilePath,
	updateState,
	type WorktreeEntry,
} from "../src/state.ts";

function entry(overrides: Partial<WorktreeEntry> = {}): WorktreeEntry {
	return {
		slug: "feat",
		path: "/nonexistent/path",
		branch: "wt/feat",
		sourcePath: "/nonexistent/source",
		sourceBranch: "main",
		sourceSha: "0".repeat(40),
		createdAt: new Date().toISOString(),
		...overrides,
	};
}

test("state roundtrip with owner-only permissions", async () => {
	const dir = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const path = join(dir, "pi-worktree", "state.json");
	const state = await loadState(path);
	assert.deepEqual(state, { version: 1, worktrees: [] });

	addEntry(state, entry());
	await saveState(state, path);

	const loaded = await loadState(path);
	assert.equal(loaded.worktrees.length, 1);
	assert.equal(loaded.worktrees[0]?.slug, "feat");
	assert.equal((await stat(path)).mode & 0o777, 0o600);
	assert.equal((await stat(dirname(path))).mode & 0o777, 0o700);

	removeEntry(loaded, "feat");
	assert.equal(loaded.worktrees.length, 0);
	assert.equal(findBySlug(loaded, "feat"), undefined);
});

test("corrupt or unknown-version state throws instead of overwriting", async () => {
	const dir = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const bad = join(dir, "bad.json");
	await writeFile(bad, "not json{", "utf8");
	await assert.rejects(() => loadState(bad), /损坏/u);

	const wrongVersion = join(dir, "v2.json");
	await writeFile(wrongVersion, JSON.stringify({ version: 2, worktrees: [] }), "utf8");
	await assert.rejects(() => loadState(wrongVersion), /版本不受支持|结构无效/u);
});

test("findByCwd matches nested subdir via realpath through a symlink", async () => {
	const root = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const target = join(root, "wt-dir");
	await mkdir(join(target, "sub", "deep"), { recursive: true });
	const link = join(root, "link");
	await symlink(target, link);

	const state = { version: 1 as const, worktrees: [entry({ path: await realpath(target) })] };
	const found = await findByCwd(state, join(link, "sub", "deep"));
	assert.equal(found?.slug, "feat");
	assert.equal(await findByCwd(state, root), undefined);
});

test("stateFilePath lives under the agent dir, not my-pi", () => {
	const previous = process.env.PI_CODING_AGENT_DIR;
	process.env.PI_CODING_AGENT_DIR = join(tmpdir(), "pi-worktree-agentdir-check");
	try {
		const path = stateFilePath();
		assert.equal(path, join(process.env.PI_CODING_AGENT_DIR, "pi-worktree", "state.json"));
		assert.ok(!path.includes("my-pi"));
	} finally {
		if (previous === undefined) delete process.env.PI_CODING_AGENT_DIR;
		else process.env.PI_CODING_AGENT_DIR = previous;
	}
});

test("updateState serializes concurrent mutations without losing entries", async () => {
	const dir = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const path = join(dir, "state.json");
	await Promise.all(
		Array.from({ length: 20 }, (_, i) =>
			updateState(
				(s) => addEntry(s, entry({ slug: `feat-${i}`, branch: `wt/feat-${i}`, path: `/p/${i}` })),
				path,
			),
		),
	);
	const loaded = await loadState(path);
	assert.equal(loaded.worktrees.length, 20);
	assert.equal(new Set(loaded.worktrees.map((e) => e.slug)).size, 20);
	assert.equal(existsSync(`${path}.lock`), false);
});

test("updateState breaks a stale lock older than the stale threshold", async () => {
	const dir = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const path = join(dir, "state.json");
	const lock = `${path}.lock`;
	await mkdir(dir, { recursive: true });
	await writeFile(lock, "1 stale", "utf8");
	const old = new Date(Date.now() - 60_000);
	await utimes(lock, old, old);

	await updateState((s) => addEntry(s, entry()), path);
	assert.equal((await loadState(path)).worktrees.length, 1);
	assert.equal(existsSync(lock), false);
});

test("updateState does not save when the mutator throws and releases the lock", async () => {
	const dir = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const path = join(dir, "state.json");
	await updateState((s) => addEntry(s, entry()), path);

	await assert.rejects(
		updateState(() => {
			throw new Error("boom");
		}, path),
		/boom/,
	);
	const loaded = await loadState(path);
	assert.equal(loaded.worktrees.length, 1);
	assert.equal(existsSync(`${path}.lock`), false);
});

test("updateState rejects when a fresh lock is held past the timeout", async () => {
	const dir = await mkdtemp(join(tmpdir(), "pi-worktree-state-"));
	const path = join(dir, "state.json");
	const lock = `${path}.lock`;
	await mkdir(dir, { recursive: true });
	await writeFile(lock, "1 fresh", "utf8");
	try {
		await assert.rejects(
			updateState(() => undefined, path, { timeoutMs: 300, staleMs: 60_000 }),
			/状态文件被占用/,
		);
	} finally {
		await unlink(lock).catch(() => undefined);
	}
});