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()