box/tests/test_tmux_auto_approver.py

#!/usr/bin/env python3
"""test_tmux_auto_approver.py — Unit tests for Tmux worker management and regex auto-approver."""

import json
import os
import sys
import tempfile
import time
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "bin"))

import tmux_auto_approver
from tmux_auto_approver import (
    DEFAULT_RULES,
    GUARDRAIL_PATTERNS,
    AutoApproverRunner,
    AutoApproverState,
    MatchRule,
    MatchVerdict,
    RegexApproverEngine,
    TmuxPaneInfo,
    TmuxWorkerTally,
    gather_tmux_tally,
    infer_agent_for_session,
)


class TestRegexApproverEngine(unittest.TestCase):
    def setUp(self):
        self.engine = RegexApproverEngine(DEFAULT_RULES)

    def test_muse_code_run_numbered_match(self):
        text = """
Terminal output before...
Would you like to run the following bash command?
  echo "hello"
› 1. Yes, proceed (y)
  2. No, skip (n)
"""
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertEqual(verdict.rule_id, "muse_code_run_numbered")
        self.assertEqual(verdict.key, "1")
        self.assertFalse(verdict.press_enter)

    def test_muse_code_run_yn_match(self):
        text = "Some actions\n› 1. Yes, proceed (y)\n"
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertIn("muse_code_run", verdict.rule_id)
        self.assertEqual(verdict.key, "1")

    def test_muse_code_allow_execution_match(self):
        text = "Allow execution of script /tmp/test.sh? [y/N]"
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertEqual(verdict.rule_id, "muse_code_allow_execution")
        self.assertEqual(verdict.key, "y")
        self.assertTrue(verdict.press_enter)

    def test_choice_abc_match(self):
        text = "Select option from the choices below:\nA) Deploy to production\nB) Staging"
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertEqual(verdict.rule_id, "choice_abc")
        self.assertEqual(verdict.key, "A")
        self.assertTrue(verdict.press_enter)

    def test_menu_numbered_match(self):
        text = "Option: Please pick a number:\n(1) Accept terms\n(2) Decline"
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertEqual(verdict.rule_id, "menu_numbered")
        self.assertEqual(verdict.key, "1")
        self.assertTrue(verdict.press_enter)

    def test_confirm_yn_match(self):
        text = "Do you want to continue? [y/n]: "
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertEqual(verdict.rule_id, "confirm_yn")
        self.assertEqual(verdict.key, "y")
        self.assertTrue(verdict.press_enter)

    def test_enter_to_continue_match(self):
        text = "Task completed successfully. Press [Enter] to continue..."
        verdict = self.engine.evaluate(text)
        self.assertTrue(verdict.matched)
        self.assertEqual(verdict.rule_id, "enter_to_continue")
        self.assertEqual(verdict.key, "Enter")

    def test_guardrail_sudo_password_blocked(self):
        text = "[sudo] password for super: "
        verdict = self.engine.evaluate(text)
        self.assertFalse(verdict.matched)
        self.assertTrue(verdict.is_blocked)
        self.assertIn("guardrail", verdict.blocked_reason.lower())

    def test_guardrail_passkey_blocked(self):
        text = "Insert security key or enter passkey PIN:"
        verdict = self.engine.evaluate(text)
        self.assertFalse(verdict.matched)
        self.assertTrue(verdict.is_blocked)
        self.assertIn("guardrail", verdict.blocked_reason.lower())

    def test_guardrail_destructive_command_blocked(self):
        text = "Running dangerous cleanup: rm -rf /"
        verdict = self.engine.evaluate(text)
        self.assertFalse(verdict.matched)
        self.assertTrue(verdict.is_blocked)
        self.assertIn("guardrail", verdict.blocked_reason.lower())

    def test_empty_text_returns_unmatched(self):
        verdict = self.engine.evaluate("")
        self.assertFalse(verdict.matched)
        self.assertFalse(verdict.is_blocked)


class TestAgentInference(unittest.TestCase):
    def test_infer_agent_names(self):
        self.assertEqual(infer_agent_for_session("/tmp/tmux-muse.sock", "worker-1"), "muse")
        self.assertEqual(infer_agent_for_session("/tmp/tmux-pip.sock", "agent-job"), "pip")
        self.assertEqual(infer_agent_for_session("/tmp/tmux-1000/default", "646-retention"), "646")
        self.assertEqual(infer_agent_for_session("/tmp/tmux-1000/default", "opm"), "opm")
        self.assertEqual(infer_agent_for_session("/tmp/tmux-1000/default", "dev-test"), "dev")
        self.assertEqual(infer_agent_for_session("/tmp/tmux-1000/default", "def"), "def")
        self.assertEqual(infer_agent_for_session("/tmp/tmux-1000/lte", "main"), "host")


class TestAutoApproverRunner(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.orig_state = tmux_auto_approver.STATE_FILE
        self.orig_audit = tmux_auto_approver.AUDIT_LOG_FILE
        tmux_auto_approver.STATE_FILE = Path(self.temp_dir.name) / "test_state.json"
        tmux_auto_approver.AUDIT_LOG_FILE = Path(self.temp_dir.name) / "test_audit.jsonl"

    def tearDown(self):
        tmux_auto_approver.STATE_FILE = self.orig_state
        tmux_auto_approver.AUDIT_LOG_FILE = self.orig_audit
        self.temp_dir.cleanup()

    @patch("tmux_auto_approver.run_tmux_cmd")
    @patch("tmux_auto_approver.capture_pane_text")
    @patch("tmux_auto_approver.gather_tmux_tally")
    def test_run_once_dry_run_dispatches_match(self, mock_tally, mock_capture, mock_tmux_cmd):
        mock_pane = TmuxPaneInfo(
            socket="/tmp/tmux-1000/default",
            session="muse",
            window_idx=0,
            pane_id="%37",
            pane_pid=12345,
            current_command="muse-bin",
            active=True,
            attached=True,
            title="muse terminal",
            agent_node="muse",
            auto_approve=True,
        )

        mock_tally.return_value = TmuxWorkerTally(
            total_sockets=1,
            total_sessions=1,
            total_panes=1,
            active_workers=1,
            by_agent={"muse": {"sessions": 1, "panes": 1, "active_commands": ["muse-bin"], "auto_approve": True}},
            panes=[mock_pane],
        )

        mock_capture.return_value = "Would you like to run the following?\n› 1. Yes, proceed (y)"

        runner = AutoApproverRunner(dry_run=True)
        results = runner.run_once()

        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["action"], "DRY_RUN_MATCH")
        self.assertEqual(results[0]["key_sent"], "1")
        self.assertEqual(results[0]["rule_name"], "Muse Code Run (Numbered)")

        # Verify deduplication within 15 seconds
        results2 = runner.run_once()
        self.assertEqual(len(results2), 0)

    @patch("tmux_auto_approver.gather_tmux_tally")
    def test_run_once_global_disabled(self, mock_tally):
        runner = AutoApproverRunner(dry_run=True)
        runner.state.global_enabled = False
        runner.state.save()
        results = runner.run_once()
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["status"], "disabled")


class TestTallyGathering(unittest.TestCase):
    @patch("tmux_auto_approver.get_existing_sockets")
    @patch("tmux_auto_approver.run_tmux_cmd")
    def test_gather_tmux_tally_parsing(self, mock_run_cmd, mock_sockets):
        mock_sockets.return_value = ["/tmp/tmux-1000/default"]
        # Format: #{session_name}___#{window_index}___#{pane_id}___#{pane_pid}___#{pane_current_command}___#{pane_active}___#{session_attached}___#{pane_title}
        mock_run_cmd.return_value = (
            0,
            "muse___0___%1___1001___muse-bin___1___1___muse worker\n"
            "dev___0___%2___1002___bash___0___0___dev terminal\n",
            "",
        )

        state = AutoApproverState()
        tally = gather_tmux_tally(state)

        self.assertEqual(tally.total_sockets, 1)
        self.assertEqual(tally.total_sessions, 2)
        self.assertEqual(tally.total_panes, 2)
        self.assertEqual(tally.active_workers, 1)  # only muse-bin is active, bash is shell
        self.assertEqual(len(tally.panes), 2)
        self.assertEqual(tally.panes[0].agent_node, "muse")
        self.assertEqual(tally.panes[1].agent_node, "dev")


class TestMuseDeferral(unittest.TestCase):
    """tmux approver must defer muse panes owned by muse_choice_watcher.

    Live double-answer regression: both daemons answered the same
    Would-you-like prompt within the same second (box-ctl + tmux audit
    overlap on %40/%0/%2), producing '11' + stray keys in the input box.
    When a per-pane muse watcher is alive, tmux must skip the pane.
    """

    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.orig_state = tmux_auto_approver.STATE_FILE
        self.orig_audit = tmux_auto_approver.AUDIT_LOG_FILE
        tmux_auto_approver.STATE_FILE = Path(self.temp_dir.name) / "test_state.json"
        tmux_auto_approver.AUDIT_LOG_FILE = Path(self.temp_dir.name) / "test_audit.jsonl"

    def tearDown(self):
        tmux_auto_approver.STATE_FILE = self.orig_state
        tmux_auto_approver.AUDIT_LOG_FILE = self.orig_audit
        self.temp_dir.cleanup()

    def _muse_pane(self, pane_id="%37", socket="/tmp/tmux-1000/default"):
        return TmuxPaneInfo(
            socket=socket, session="muse", window_idx=0, pane_id=pane_id,
            pane_pid=12345, current_command="muse-bin", active=True,
            attached=True, title="muse terminal", agent_node="muse",
            auto_approve=True,
        )

    def _tally(self, panes):
        return TmuxWorkerTally(
            total_sockets=1, total_sessions=1, total_panes=len(panes),
            active_workers=len(panes), by_agent={}, panes=panes,
        )

    @patch("tmux_auto_approver.run_tmux_cmd")
    @patch("tmux_auto_approver.capture_pane_text")
    @patch("tmux_auto_approver.gather_tmux_tally")
    def test_muse_pane_skipped_when_watcher_alive(
            self, mock_tally, mock_capture, mock_tmux_cmd):
        import muse_choice_watcher as mcw
        mock_tally.return_value = self._tally([self._muse_pane()])
        mock_capture.return_value = (
            "Would you like to run the following?\n› 1. Yes, proceed (y)")
        with patch.object(mcw, "is_running", return_value=99999):
            runner = AutoApproverRunner(dry_run=True)
            results = runner.run_once()
        self.assertEqual(results, [])
        mock_tmux_cmd.assert_not_called()

    @patch("tmux_auto_approver.run_tmux_cmd")
    @patch("tmux_auto_approver.capture_pane_text")
    @patch("tmux_auto_approver.gather_tmux_tally")
    def test_non_muse_pane_still_approved(
            self, mock_tally, mock_capture, mock_tmux_cmd):
        pane = TmuxPaneInfo(
            socket="/tmp/tmux-pip.sock", session="worker", window_idx=0,
            pane_id="%1", pane_pid=999, current_command="agent-worker",
            active=True, attached=True, title="w", agent_node="pip",
            auto_approve=True,
        )
        mock_tally.return_value = self._tally([pane])
        mock_capture.return_value = (
            "Would you like to run the following?\n› 1. Yes, proceed (y)")
        runner = AutoApproverRunner(dry_run=True)
        results = runner.run_once()
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["action"], "DRY_RUN_MATCH")


class TestDedupSocketScoped(unittest.TestCase):
    """Dedup must be keyed by socket:pane, not bare pane id.

    Same %N exists on every tmux socket; bare-pane dedup suppresses a
    real prompt on socket B because socket A saw one (the %0-on-two-
    sockets collision, tmux-side).
    """

    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.orig_state = tmux_auto_approver.STATE_FILE
        self.orig_audit = tmux_auto_approver.AUDIT_LOG_FILE
        tmux_auto_approver.STATE_FILE = Path(self.temp_dir.name) / "test_state.json"
        tmux_auto_approver.AUDIT_LOG_FILE = Path(self.temp_dir.name) / "test_audit.jsonl"

    def tearDown(self):
        tmux_auto_approver.STATE_FILE = self.orig_state
        tmux_auto_approver.AUDIT_LOG_FILE = self.orig_audit
        self.temp_dir.cleanup()

    @patch("tmux_auto_approver.run_tmux_cmd")
    @patch("tmux_auto_approver.capture_pane_text")
    @patch("tmux_auto_approver.gather_tmux_tally")
    def test_same_pane_id_on_two_sockets_both_approved(
            self, mock_tally, mock_capture, mock_tmux_cmd):
        def mk(sock):
            return TmuxPaneInfo(
                socket=sock, session="w", window_idx=0, pane_id="%1",
                pane_pid=999, current_command="agent-worker", active=True,
                attached=True, title="w", agent_node="pip",
                auto_approve=True,
            )
        mock_tally.return_value = TmuxWorkerTally(
            total_sockets=2, total_sessions=2, total_panes=2,
            active_workers=2, by_agent={},
            panes=[mk("/tmp/tmux-pip.sock"), mk("/tmp/tmux-opm.sock")],
        )
        mock_capture.return_value = (
            "Would you like to run the following?\n› 1. Yes, proceed (y)")
        runner = AutoApproverRunner(dry_run=True)
        results = runner.run_once()
        self.assertEqual(len(results), 2)


class TestCaptureJoinWrapped(unittest.TestCase):
    def test_capture_joins_wrapped_lines(self):
        with patch("tmux_auto_approver.run_tmux_cmd",
                   return_value=(0, "ok", "")) as m:
            tmux_auto_approver.capture_pane_text("/tmp/s", "%1", lines=30)
        args = m.call_args[0]
        self.assertIn("-J", args)


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