90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
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()
|