amc/tests/test_control.py

import io
import json
import os
import subprocess
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

import amc_server.mixins.control as control
from amc_server.mixins.control import SessionControlMixin


class DummyControlHandler(SessionControlMixin):
    def __init__(self, body=None):
        if body is None:
            body = {}
        raw = json.dumps(body).encode("utf-8")
        self.headers = {"Content-Length": str(len(raw))}
        self.rfile = io.BytesIO(raw)
        self.sent = []
        self.errors = []

    def _send_json(self, code, payload):
        self.sent.append((code, payload))

    def _json_error(self, code, message):
        self.errors.append((code, message))


class SessionControlMixinTests(unittest.TestCase):
    def _write_session(self, sessions_dir: Path, session_id: str, zellij_session="infra", zellij_pane="21"):
        sessions_dir.mkdir(parents=True, exist_ok=True)
        session_file = sessions_dir / f"{session_id}.json"
        session_file.write_text(
            json.dumps({
                "session_id": session_id,
                "zellij_session": zellij_session,
                "zellij_pane": zellij_pane,
            })
        )

    def test_inject_text_then_enter_is_two_step_with_delay(self):
        handler = DummyControlHandler()
        calls = []

        def fake_inject(zellij_session, pane_id, text, send_enter=True):
            calls.append((zellij_session, pane_id, text, send_enter))
            return {"ok": True}

        handler._inject_to_pane = fake_inject

        with patch("amc_server.mixins.control.time.sleep") as sleep_mock, patch.dict(os.environ, {}, clear=True):
            result = handler._inject_text_then_enter("infra", 24, "testing")

        self.assertEqual(result, {"ok": True})
        self.assertEqual(
            calls,
            [
                ("infra", 24, "testing", False),
                ("infra", 24, "", True),
            ],
        )
        sleep_mock.assert_called_once_with(0.20)

    def test_inject_text_then_enter_delay_honors_environment_override(self):
        handler = DummyControlHandler()
        handler._inject_to_pane = MagicMock(return_value={"ok": True})

        with patch("amc_server.mixins.control.time.sleep") as sleep_mock, patch.dict(
            os.environ, {"AMC_SUBMIT_ENTER_DELAY_MS": "350"}, clear=True
        ):
            result = handler._inject_text_then_enter("infra", 9, "hello")

        self.assertEqual(result, {"ok": True})
        sleep_mock.assert_called_once_with(0.35)

    def test_respond_to_session_freeform_selects_other_then_submits_text(self):
        session_id = "abc123"
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir)
            self._write_session(sessions_dir, session_id)

            handler = DummyControlHandler(
                {
                    "text": "testing",
                    "freeform": True,
                    "optionCount": 3,
                }
            )

            with patch.object(control, "SESSIONS_DIR", sessions_dir), patch(
                "amc_server.mixins.control.time.sleep"
            ) as sleep_mock:
                handler._inject_to_pane = MagicMock(return_value={"ok": True})
                handler._inject_text_then_enter = MagicMock(return_value={"ok": True})
                handler._respond_to_session(session_id)

            handler._inject_to_pane.assert_called_once_with("infra", 21, "4", send_enter=False)
            handler._inject_text_then_enter.assert_called_once_with("infra", 21, "testing")
            sleep_mock.assert_called_once_with(handler._FREEFORM_MODE_SWITCH_DELAY_SEC)
            self.assertEqual(handler.errors, [])
            self.assertEqual(handler.sent, [(200, {"ok": True})])

    def test_respond_to_session_non_freeform_uses_text_then_enter_helper(self):
        session_id = "abc456"
        with tempfile.TemporaryDirectory() as tmpdir:
            sessions_dir = Path(tmpdir)
            self._write_session(sessions_dir, session_id, zellij_pane="terminal_5")

            handler = DummyControlHandler({"text": "hello"})

            with patch.object(control, "SESSIONS_DIR", sessions_dir):
                handler._inject_to_pane = MagicMock(return_value={"ok": True})
                handler._inject_text_then_enter = MagicMock(return_value={"ok": True})
                handler._respond_to_session(session_id)

            handler._inject_to_pane.assert_not_called()
            handler._inject_text_then_enter.assert_called_once_with("infra", 5, "hello")
            self.assertEqual(handler.sent, [(200, {"ok": True})])

    def test_respond_to_session_missing_session_returns_404(self):
        handler = DummyControlHandler({"text": "hello"})
        sessions_dir = Path(self.id())

        with patch.object(control, "SESSIONS_DIR", sessions_dir):
            handler._respond_to_session("does-not-exist")

        self.assertEqual(handler.sent, [])
        self.assertEqual(handler.errors, [(404, "Session not found")])

    def test_try_plugin_inject_uses_explicit_session_and_payload(self):
        handler = DummyControlHandler()
        env = {}
        completed = subprocess.CompletedProcess(args=[], returncode=0, stdout="", stderr="")

        with patch.object(control, "ZELLIJ_BIN", "/opt/homebrew/bin/zellij"), patch(
            "amc_server.mixins.control.subprocess.run", return_value=completed
        ) as run_mock:
            result = handler._try_plugin_inject(env, "infra", 24, "testing", send_enter=True)

        self.assertEqual(result, {"ok": True})
        args = run_mock.call_args.args[0]
        self.assertEqual(args[0], "/opt/homebrew/bin/zellij")
        self.assertEqual(args[1:3], ["--session", "infra"])
        payload = json.loads(args[-1])
        self.assertEqual(payload["pane_id"], 24)
        self.assertEqual(payload["text"], "testing")
        self.assertIs(payload["send_enter"], True)

    def test_inject_to_pane_without_plugin_or_unsafe_fallback_returns_error(self):
        handler = DummyControlHandler()

        with patch.object(control, "ZELLIJ_PLUGIN", Path("/definitely/missing/plugin.wasm")), patch.dict(
            os.environ, {}, clear=True
        ):
            result = handler._inject_to_pane("infra", 24, "testing", send_enter=True)

        self.assertFalse(result["ok"])
        self.assertIn("Pane-targeted injection requires zellij-send-keys plugin", result["error"])

    def test_inject_to_pane_allows_unsafe_fallback_when_enabled(self):
        handler = DummyControlHandler()

        with patch.object(control, "ZELLIJ_PLUGIN", Path("/definitely/missing/plugin.wasm")), patch.dict(
            os.environ, {"AMC_ALLOW_UNSAFE_WRITE_CHARS_FALLBACK": "1"}, clear=True
        ):
            handler._try_write_chars_inject = MagicMock(return_value={"ok": True})
            result = handler._inject_to_pane("infra", 24, "testing", send_enter=True)

        self.assertEqual(result, {"ok": True})
        handler._try_write_chars_inject.assert_called_once()


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