DocForge2/tests/test_command_reference_tool.py

from __future__ import annotations

import asyncio
import contextlib
import hashlib
import io
import os
import tempfile
import unittest
from pathlib import Path

from docforge.mcp_server import ALL_TOOLS, APPLICATION_TOOLS
from tools.generate_command_reference import (
    EXPECTED_CLI_ROWS,
    EXPECTED_MCP_ROWS,
    CommandReferenceDrift,
    CommandReferenceToolError,
    collect_command_reference_rows,
    generate_command_reference_bytes,
    main,
    publish_or_check_command_reference,
)

ROOT = Path(__file__).resolve().parents[1]
ALPHA = ROOT / "tests" / "fixtures" / "alpha"


def _tree_snapshot(root: Path) -> tuple[tuple[str, str, str], ...]:
    records: list[tuple[str, str, str]] = []
    for path in sorted(root.rglob("*")):
        relative = path.relative_to(root).as_posix()
        if path.is_symlink():
            records.append((relative, "symlink", os.readlink(path)))
        elif path.is_dir():
            records.append((relative, "directory", ""))
        else:
            records.append(
                (
                    relative,
                    "file",
                    hashlib.sha256(path.read_bytes()).hexdigest(),
                )
            )
    return tuple(records)


class CommandReferenceToolTests(unittest.TestCase):
    def test_generated_bytes_are_stable_and_never_mutate_the_fixture(self) -> None:
        before = _tree_snapshot(ALPHA)

        first = generate_command_reference_bytes(ALPHA)
        middle = _tree_snapshot(ALPHA)
        second = generate_command_reference_bytes(ALPHA)

        self.assertEqual(first, second)
        self.assertEqual(before, middle)
        self.assertEqual(before, _tree_snapshot(ALPHA))
        self.assertTrue(first.endswith(b"\n"))
        self.assertEqual(1, first.count(b"## CLI commands"))
        self.assertEqual(1, first.count(b"## MCP tools"))

    def test_registration_collection_is_complete(self) -> None:
        cli_rows, mcp_rows = asyncio.run(
            collect_command_reference_rows(
                ALPHA,
                proposal_writer="alpha-editor",
                canonical_applier="alpha-editor",
            )
        )

        self.assertEqual(EXPECTED_CLI_ROWS, len(cli_rows))
        self.assertEqual(EXPECTED_MCP_ROWS, len(mcp_rows))
        self.assertEqual(len(cli_rows), len({item.name for item in cli_rows}))
        self.assertEqual(
            set((*ALL_TOOLS, *APPLICATION_TOOLS)),
            {item.name for item in mcp_rows},
        )
        self.assertEqual(
            {"read", "proposal", "application"},
            {item.surface for item in mcp_rows},
        )

    def test_atomic_write_check_and_drift_detection(self) -> None:
        content = generate_command_reference_bytes(ALPHA)
        with tempfile.TemporaryDirectory() as directory:
            repository = Path(directory).resolve()
            (repository / "generated").mkdir()
            output = Path("generated/command-reference.md")
            target = repository / output

            self.assertEqual(
                "written",
                publish_or_check_command_reference(
                    content,
                    repository_root=repository,
                    project_root=ALPHA,
                    output=output,
                    check=False,
                ),
            )
            self.assertEqual(content, target.read_bytes())
            self.assertEqual(
                "current",
                publish_or_check_command_reference(
                    content,
                    repository_root=repository,
                    project_root=ALPHA,
                    output=output,
                    check=True,
                ),
            )
            target.write_text("stale\n", encoding="utf-8")
            with self.assertRaises(CommandReferenceDrift):
                publish_or_check_command_reference(
                    content,
                    repository_root=repository,
                    project_root=ALPHA,
                    output=output,
                    check=True,
                )
            self.assertEqual(b"stale\n", target.read_bytes())
            self.assertEqual(
                "written",
                publish_or_check_command_reference(
                    content,
                    repository_root=repository,
                    project_root=ALPHA,
                    output=output,
                    check=False,
                ),
            )
            self.assertEqual(content, target.read_bytes())
            self.assertEqual(
                "unchanged",
                publish_or_check_command_reference(
                    content,
                    repository_root=repository,
                    project_root=ALPHA,
                    output=output,
                    check=False,
                ),
            )

    def test_output_confinement_rejects_escapes_and_symlinks(self) -> None:
        content = b"reference\n"
        with tempfile.TemporaryDirectory() as directory:
            repository = Path(directory).resolve()
            generated = repository / "generated"
            generated.mkdir()
            outside = repository.parent / f"{repository.name}-outside.md"
            symlink = generated / "reference.md"
            symlink.symlink_to(outside)

            for output in (Path("../escape.md"), Path("generated/reference.md")):
                with (
                    self.subTest(output=output),
                    self.assertRaises(CommandReferenceToolError),
                ):
                    publish_or_check_command_reference(
                        content,
                        repository_root=repository,
                        project_root=ALPHA,
                        output=output,
                        check=False,
                    )
            symlink.unlink()
            (repository / "linked").symlink_to(generated, target_is_directory=True)
            with self.assertRaises(CommandReferenceToolError):
                publish_or_check_command_reference(
                    content,
                    repository_root=repository,
                    project_root=ALPHA,
                    output=Path("linked/reference.md"),
                    check=False,
                )

        with self.assertRaises(CommandReferenceToolError):
            publish_or_check_command_reference(
                content,
                repository_root=ROOT,
                project_root=ALPHA,
                output=Path("tests/fixtures/alpha/generated-reference.md"),
                check=False,
            )

    def test_command_runs_with_explicit_repository_and_project_roots(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            repository = Path(directory).resolve()
            (repository / "generated").mkdir()
            output = Path("generated/reference.md")
            stdout = io.StringIO()

            with contextlib.redirect_stdout(stdout):
                status = main(
                    (
                        "--repository-root",
                        str(repository),
                        "--project-root",
                        str(ALPHA),
                        "--output",
                        str(output),
                    )
                )

            self.assertEqual(0, status)
            self.assertTrue((repository / output).is_file())
            self.assertIn('"cli_rows":28', stdout.getvalue())
            self.assertIn('"mcp_rows":36', stdout.getvalue())

            target = repository / output
            target.write_text("drift\n", encoding="utf-8")
            stderr = io.StringIO()
            with contextlib.redirect_stderr(stderr):
                check_status = main(
                    (
                        "--repository-root",
                        str(repository),
                        "--project-root",
                        str(ALPHA),
                        "--output",
                        str(output),
                        "--check",
                    )
                )

            self.assertEqual(1, check_status)
            self.assertEqual("drift\n", target.read_text(encoding="utf-8"))
            self.assertIn("missing or stale", stderr.getvalue())


if __name__ == "__main__":
    unittest.main()