tai/tests/bashrc.py

#!/usr/bin/env python3
"""tests/bashrc.py — the tests that can only run honestly inside the throwaway
container (see tests/docker/Dockerfile, docker-compose.yml, tests/container.py).

Everything here either installs a real hook into a real $HOME/.bashrc and
spawns a real `bash -i` to watch it fire, or drives a real `tmux` session.
Both are things no test should ever do against the machine actually running
Claude Code -- so every test case that touches $HOME/.bashrc checks
TAI_CONTAINER_TEST=1 in setup() and refuses outright if it's missing, rather
than trusting whoever runs this file to remember to only run it in the
container. tests/docker/Dockerfile sets that var; nothing else should.

Deterministic throughout, same as tests/plexer.py: no LLM involved, this is
pty/shell/signal mechanics, not a model decision.
"""
import asyncio
import glob
import os
import re
import shutil
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from tai_test import TestCase, TestAssertionFailed, ManagedProcess, main_cli  # noqa: E402
import tai  # noqa: E402


async def _drain_forever(proc):
    """Same requirement as tests/plexer.py's helper of the same name (see
    that file for the full postmortem): only needed when the spawned process
    itself runs a second relay loop that would otherwise self-deadlock on its
    own unread output -- i.e. once bash execs into `tai.py --plexer-wrap`.
    Plain `bash -i` that never wraps doesn't need this; those tests read
    directly instead."""
    try:
        while True:
            chunk = await proc._read_chunk()
            if not chunk:
                await asyncio.sleep(0.05)
    except asyncio.CancelledError:
        pass


class _BashrcHookTestBase(TestCase):
    """Shared fixture for every test that installs the real plexer hook into
    a real .bashrc: a throwaway $HOME, the real tai.py binary copied into
    $HOME/.local/bin via tai.op_binary (exactly what `tai.py --install` does
    on a real machine), and the hook itself via tai.op_plexer_hook -- the
    same functions a real install runs, not a hand-rolled substitute."""

    mode = "pty"
    timeout = 30.0

    async def setup(self):
        if os.environ.get("TAI_CONTAINER_TEST") != "1":
            raise TestAssertionFailed(
                "refusing to run (TAI_CONTAINER_TEST is not set): this test installs a real "
                "hook into $HOME/.bashrc and must only ever run inside the disposable test "
                "container (see tests/docker/Dockerfile), never against a real machine's home"
            )
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        os.makedirs(os.path.join(self.tmp_home, ".local", "bin"), exist_ok=True)
        real_home = os.environ.get("HOME")
        os.environ["HOME"] = self.tmp_home
        try:
            tai.op_binary("install", self.tmp_home)
            tai.op_plexer_hook("install", self.tmp_home)
        finally:
            if real_home is not None:
                os.environ["HOME"] = real_home
            else:
                os.environ.pop("HOME", None)
        self._drain_task = None

    async def _spawn_bash(self, args, mode=None, env_overrides=None):
        env = self.build_env()
        env.update(env_overrides or {})
        proc = ManagedProcess(["bash"] + list(args), mode=mode or self.mode, env=env, on_line=self._on_line)
        await proc.start()
        self._processes.append(proc)
        return proc

    def _sock_matches(self):
        return glob.glob(os.path.join(self.tmp_home, ".tai", "plexer", "*.sock"))

    async def teardown(self):
        if self._drain_task is not None:
            self._drain_task.cancel()
        shutil.rmtree(self.tmp_home, ignore_errors=True)


class BashrcInstallAndAutoWrapTest(_BashrcHookTestBase):
    """The actual point of building a container for this: a real `bash -i`,
    with a real installed hook in its real $HOME/.bashrc, auto-wraps into a
    real NativePlexerBackend with no human and no flag telling it to."""

    name = "bashrc-install-and-auto-wrap"

    async def run(self):
        proc = await self._spawn_bash(["-i"], env_overrides={"HOME": self.tmp_home})
        self._drain_task = asyncio.ensure_future(_drain_forever(proc))

        sock_path = None
        for _ in range(50):
            matches = self._sock_matches()
            if matches:
                sock_path = matches[0]
                break
            await asyncio.sleep(0.2)
        if not sock_path:
            raise TestAssertionFailed(
                "bash -i never auto-wrapped into a plexer backend via the installed .bashrc hook"
            )

        controller = tai.NativePlexerController(sock_path)
        await asyncio.sleep(0.3)
        status = controller.status()

        controller.send_keys("echo $TAI_PLEXER_SOCKET", enter=True)
        await asyncio.sleep(0.4)
        pane = controller.read_pane(50)

        return {"sock_path": sock_path, "pane": pane, "status": status}

    async def verify(self, result):
        if result["sock_path"] not in result["pane"]:
            raise TestAssertionFailed(
                "the auto-wrapped inner shell's $TAI_PLEXER_SOCKET did not match the backend's "
                "own socket path: %r" % result["pane"][-300:]
            )
        if not result["status"].get("pid") or result["status"]["pid"] <= 0:
            raise TestAssertionFailed("status() did not report a live child pid after auto-wrap")


class AntiRecursionTest(_BashrcHookTestBase):
    """Already-wrapped ($TAI_PLEXER_ACTIVE=1) must never wrap again."""

    name = "bashrc-hook-anti-recursion"

    async def run(self):
        proc = await self._spawn_bash(
            ["-i"], env_overrides={"HOME": self.tmp_home, "TAI_PLEXER_ACTIVE": "1"}
        )
        await proc.send_line("echo anti_recursion_marker_ok")
        await proc.wait_for_pattern(re.compile(r"anti_recursion_marker_ok"), timeout=8.0)
        await asyncio.sleep(0.3)
        return {"sock_matches": self._sock_matches()}

    async def verify(self, result):
        if result["sock_matches"]:
            raise TestAssertionFailed(
                "a plexer socket was created even though TAI_PLEXER_ACTIVE was already set -- "
                "recursive wrap was not prevented"
            )


class TmuxActiveSkipTest(_BashrcHookTestBase):
    """An already-active tmux pane ($TMUX set) must also never wrap."""

    name = "bashrc-hook-tmux-active-skip"

    async def run(self):
        proc = await self._spawn_bash(
            ["-i"], env_overrides={"HOME": self.tmp_home, "TMUX": "dummy-session:0.0"}
        )
        await proc.send_line("echo tmux_skip_marker_ok")
        await proc.wait_for_pattern(re.compile(r"tmux_skip_marker_ok"), timeout=8.0)
        await asyncio.sleep(0.3)
        return {"sock_matches": self._sock_matches()}

    async def verify(self, result):
        if result["sock_matches"]:
            raise TestAssertionFailed(
                "a plexer socket was created even though $TMUX was already set -- tmux should "
                "have been left in sole control of the session"
            )


class NonInteractiveSkipTest(_BashrcHookTestBase):
    """A non-interactive shell must never wrap, even if it explicitly sources
    the hook (isolates the hook's own `case $- in *i*)` guard from bash's
    separate, unrelated default of not reading ~/.bashrc for `-c` at all)."""

    name = "bashrc-hook-non-interactive-skip"
    mode = "pipe"

    async def run(self):
        proc = await self._spawn_bash(
            ["-c", 'source "$HOME/.bashrc" 2>&1; echo marker_non_interactive_ok'],
            mode="pipe",
            env_overrides={"HOME": self.tmp_home},
        )
        output = await proc.drain_until_eof(timeout=10.0)
        returncode = await proc.wait(timeout=5.0)
        return {"output": output, "returncode": returncode}

    async def verify(self, result):
        if "marker_non_interactive_ok" not in result["output"]:
            raise TestAssertionFailed(
                "a non-interactive `bash -c` that explicitly sourced the hook never reached past "
                "it: %r" % result["output"][-500:]
            )
        if result["returncode"] != 0:
            raise TestAssertionFailed(
                "non-interactive shell exited %r after sourcing the hook" % result["returncode"]
            )
        if self._sock_matches():
            raise TestAssertionFailed(
                "a plexer socket was created from a non-interactive shell -- the "
                "`case $- in *i*)` guard failed"
            )


class TmuxControllerLiveTest(TestCase):
    """Drives TmuxPlexerController against a real `tmux new-session -d`, not
    a mock -- the one TmuxPlexerController code path tests/plexer.py cannot
    cover, since that file only ever exercises NativePlexerBackend."""

    name = "bashrc-tmux-controller-live"
    timeout = 20.0

    async def setup(self):
        if not shutil.which("tmux"):
            raise TestAssertionFailed(
                "tmux is not installed -- required to test TmuxPlexerController against a real session"
            )
        self.session_name = "tai-test-%d" % os.getpid()

    async def run(self):
        create = await asyncio.create_subprocess_exec(
            "tmux", "new-session", "-d", "-s", self.session_name, "-x", "80", "-y", "24",
        )
        await create.wait()
        if create.returncode != 0:
            raise TestAssertionFailed("tmux new-session -d failed (returncode %r)" % create.returncode)

        pane = "%s:0.0" % self.session_name
        controller = tai.TmuxPlexerController(pane)

        status = controller.status()
        if status.get("backend") != "tmux":
            raise TestAssertionFailed("status() did not report backend=tmux: %r" % status)

        controller.send_keys("echo tmux_controller_marker_ok", enter=True)
        await asyncio.sleep(0.6)
        pane_text = controller.read_pane(50)
        if "tmux_controller_marker_ok" not in pane_text:
            raise TestAssertionFailed(
                "send_keys/read_pane against a real tmux session failed: %r" % pane_text[-300:]
            )

        controller.type_text("echo tmux_typed_marker_ok", delay=0.01, enter=True)
        await asyncio.sleep(0.6)
        pane_text_2 = controller.read_pane(50)
        if "tmux_typed_marker_ok" not in pane_text_2:
            raise TestAssertionFailed(
                "type_text against a real tmux session failed: %r" % pane_text_2[-300:]
            )

        bogus = tai.TmuxPlexerController("nonexistent-session-%d:99.99" % os.getpid())
        raised = False
        try:
            bogus.status()
        except tai.PlexerError:
            raised = True
        return {"raised": raised}

    async def verify(self, result):
        if not result["raised"]:
            raise TestAssertionFailed(
                "TmuxPlexerController against a nonexistent pane should raise PlexerError, not "
                "silently succeed"
            )

    async def teardown(self):
        kill = await asyncio.create_subprocess_exec(
            "tmux", "kill-session", "-t", self.session_name,
            stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL,
        )
        await kill.wait()


class PlexerHookInstallUninstallRoundTripTest(TestCase):
    """Purely deterministic, no subprocess, no real $HOME -- the install
    functions are tested directly against a throwaway file path, so this one
    is safe to run anywhere, not just in the container."""

    name = "bashrc-hook-install-uninstall-roundtrip"
    timeout = 10.0

    async def setup(self):
        self.tmp_dir = tempfile.mkdtemp(prefix="tai-test-bashrc-")
        self.bashrc_path = os.path.join(self.tmp_dir, ".bashrc")
        with open(self.bashrc_path, "w", encoding="utf-8") as handle:
            handle.write("# pre-existing user content\nexport FOO=bar\n")

    async def run(self):
        with open(self.bashrc_path, encoding="utf-8") as handle:
            original = handle.read()

        installed_first = tai.upsert_plexer_bashrc_block(self.bashrc_path)
        present_after_first = tai.plexer_hook_present(self.bashrc_path)
        with open(self.bashrc_path, encoding="utf-8") as handle:
            content_after_first = handle.read()

        installed_second = tai.upsert_plexer_bashrc_block(self.bashrc_path)
        with open(self.bashrc_path, encoding="utf-8") as handle:
            content_after_second = handle.read()

        backup_path = self.bashrc_path + ".bak-tai"
        backup_exists = os.path.exists(backup_path)
        backup_content = None
        if backup_exists:
            with open(backup_path, encoding="utf-8") as handle:
                backup_content = handle.read()

        removed = tai.remove_plexer_bashrc_block(self.bashrc_path)
        present_after_remove = tai.plexer_hook_present(self.bashrc_path)
        with open(self.bashrc_path, encoding="utf-8") as handle:
            content_after_remove = handle.read()

        reinstalled = tai.upsert_plexer_bashrc_block(self.bashrc_path)
        present_after_reinstall = tai.plexer_hook_present(self.bashrc_path)

        return {
            "original": original,
            "installed_first": installed_first,
            "present_after_first": present_after_first,
            "content_after_first": content_after_first,
            "installed_second": installed_second,
            "content_after_second": content_after_second,
            "backup_exists": backup_exists,
            "backup_content": backup_content,
            "removed": removed,
            "present_after_remove": present_after_remove,
            "content_after_remove": content_after_remove,
            "reinstalled": reinstalled,
            "present_after_reinstall": present_after_reinstall,
        }

    async def verify(self, result):
        if not result["installed_first"]:
            raise TestAssertionFailed("first upsert_plexer_bashrc_block() reported no change on a clean file")
        if not result["present_after_first"]:
            raise TestAssertionFailed("plexer_hook_present() returned false right after installing")
        if "export FOO=bar" not in result["content_after_first"]:
            raise TestAssertionFailed("installing the hook clobbered pre-existing .bashrc content")

        if result["installed_second"]:
            raise TestAssertionFailed("upsert_plexer_bashrc_block() was not idempotent -- reported a change on a second call")
        if result["content_after_second"] != result["content_after_first"]:
            raise TestAssertionFailed("a second install duplicated the hook block instead of being a no-op")

        if not result["backup_exists"]:
            raise TestAssertionFailed("no .bak-tai backup was created before modifying .bashrc")
        if result["backup_content"] != result["original"]:
            raise TestAssertionFailed("the .bak-tai backup does not match the original pre-install content")

        if not result["removed"]:
            raise TestAssertionFailed("remove_plexer_bashrc_block() reported no change when the hook was present")
        if result["present_after_remove"]:
            raise TestAssertionFailed("plexer_hook_present() still true after removal")
        if "export FOO=bar" not in result["content_after_remove"]:
            raise TestAssertionFailed("uninstalling the hook lost pre-existing .bashrc content")
        if tai.PLEXER_BASHRC_MARK_BEGIN in result["content_after_remove"]:
            raise TestAssertionFailed("uninstall left hook marker text behind")

        if not result["reinstalled"]:
            raise TestAssertionFailed("could not reinstall the hook after a clean uninstall")
        if not result["present_after_reinstall"]:
            raise TestAssertionFailed("hook not present after reinstalling post-uninstall")

    async def teardown(self):
        shutil.rmtree(self.tmp_dir, ignore_errors=True)


ALL_CASES = [
    BashrcInstallAndAutoWrapTest(),
    AntiRecursionTest(),
    TmuxActiveSkipTest(),
    NonInteractiveSkipTest(),
    TmuxControllerLiveTest(),
    PlexerHookInstallUninstallRoundTripTest(),
]

if __name__ == "__main__":
    sys.exit(main_cli(ALL_CASES, default_log_dir=str(Path(__file__).resolve().parent / "logs")))