Initial implementation

tests/test_client.py

@@ -0,0 +1,89 @@
+import asyncio
+import signal
+import unittest
+from unittest.mock import AsyncMock, patch
+
+from rex.client import action_manager, receive_actions, run_command
+from rex.client.action_manager import ClientAction
+
+
+class FakeWebSocket:
+    def __init__(self, messages: list[str]) -> None:
+        self.messages = iter(messages)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> str:
+        await asyncio.sleep(0)
+        try:
+            return next(self.messages)
+        except StopIteration:
+            raise StopAsyncIteration
+
+
+class ClientTests(unittest.IsolatedAsyncioTestCase):
+    async def test_run_command_waits_for_process(self) -> None:
+        process = AsyncMock()
+        process.wait.return_value = 0
+
+        with patch("rex.client.asyncio.create_subprocess_shell", return_value=process):
+            await run_command("test command")
+
+        process.wait.assert_awaited_once()
+
+    async def test_cancelling_command_terminates_process_group(self) -> None:
+        process = AsyncMock()
+        process.pid = 123
+        process.returncode = None
+        wait_started = asyncio.Event()
+        wait_calls = 0
+
+        async def wait() -> int:
+            nonlocal wait_calls
+            wait_calls += 1
+            if wait_calls == 1:
+                wait_started.set()
+                await asyncio.Event().wait()
+            return 0
+
+        process.wait.side_effect = wait
+
+        with (
+            patch("rex.client.asyncio.create_subprocess_shell", return_value=process),
+            patch("rex.client.os.name", "posix"),
+            patch("rex.client.os.killpg") as killpg,
+        ):
+            task = asyncio.create_task(run_command("test command"))
+            await wait_started.wait()
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+        killpg.assert_called_once_with(123, signal.SIGTERM)
+        self.assertEqual(process.wait.await_count, 2)
+
+    async def test_disconnect_cancels_running_commands(self) -> None:
+        action = ClientAction(name="wake", command="test command")
+        command_started = asyncio.Event()
+        command_cancelled = asyncio.Event()
+
+        async def running_command(_: str) -> None:
+            command_started.set()
+            try:
+                await asyncio.Event().wait()
+            finally:
+                command_cancelled.set()
+
+        websocket = FakeWebSocket(['{"name":"wake"}'])
+        with (
+            patch.object(action_manager, "get_action", return_value=action),
+            patch("rex.client.run_command", side_effect=running_command),
+        ):
+            await receive_actions(websocket)  # type: ignore[arg-type]
+
+        self.assertTrue(command_started.is_set())
+        self.assertTrue(command_cancelled.is_set())
+
+
+if __name__ == "__main__":
+    unittest.main()