Harden and rewrite for release
This commit is contained in:
@@ -1,89 +0,0 @@
|
||||
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()
|
||||
@@ -0,0 +1,46 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from rex.client.config import ClientAction, ClientConfig, ClientConfigManager
|
||||
from rex.server.config import ServerConfigManager
|
||||
from rex.server.connection_manager import ConnectionManager
|
||||
|
||||
|
||||
class ConfigTests(unittest.TestCase):
|
||||
def test_server_config_is_toml_and_persists_scoped_keys(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "server.toml"
|
||||
manager = ServerConfigManager(path)
|
||||
key = manager.create_key("device", {"connect"})
|
||||
self.assertIsNotNone(key)
|
||||
self.assertIn("[[keys]]", path.read_text())
|
||||
self.assertIsNotNone(ServerConfigManager(path).get_key(key.secret, "connect"))
|
||||
self.assertIsNone(ServerConfigManager(path).get_key(key.secret, "execute"))
|
||||
|
||||
def test_client_url_omits_optional_port(self) -> None:
|
||||
self.assertEqual(ClientConfig().websocket_url(), "wss://127.0.0.1/ws/default-device")
|
||||
self.assertEqual(ClientConfig(endpoint="rex.test", scheme="ws", port=8080).websocket_url(), "ws://rex.test:8080/ws/default-device")
|
||||
|
||||
def test_client_actions_are_local_argv_only(self) -> None:
|
||||
action = ClientAction(name="lock", argv=["/bin/echo", "locked"])
|
||||
self.assertEqual(action.argv[0], "/bin/echo")
|
||||
with self.assertRaises(ValueError):
|
||||
ClientAction(name="bad name", argv=["/bin/echo"])
|
||||
|
||||
|
||||
class ConnectionManagerTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_only_registered_actions_are_sent(self) -> None:
|
||||
class Socket:
|
||||
def __init__(self) -> None:
|
||||
self.messages: list[dict[str, str]] = []
|
||||
async def accept(self) -> None: pass
|
||||
async def send_json(self, message: dict[str, str]) -> None: self.messages.append(message)
|
||||
|
||||
socket = Socket()
|
||||
manager = ConnectionManager()
|
||||
await manager.connect("desk", socket, {"lock"}) # type: ignore[arg-type]
|
||||
self.assertFalse(await manager.send("desk", "shell"))
|
||||
self.assertTrue(await manager.send("desk", "lock"))
|
||||
self.assertEqual(socket.messages, [{"type": "action", "name": "lock"}])
|
||||
|
||||
@@ -1,125 +0,0 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
from textual.widgets import DataTable, TabbedContent, Tree
|
||||
|
||||
from rex.management.app import RexManager
|
||||
from rex.management.client import ManagementClient, ManagementError
|
||||
from rex.server.device_manager import DeviceManager
|
||||
|
||||
|
||||
class ManagementClientTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_reads_keys_and_encodes_names(self) -> None:
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if request.method == "GET":
|
||||
return httpx.Response(
|
||||
200, json={"keys": [{"name": "home", "secret": "abc"}]}
|
||||
)
|
||||
return httpx.Response(200, json={"key": {}})
|
||||
|
||||
client = ManagementClient(
|
||||
"http://rex.test", transport=httpx.MockTransport(handler)
|
||||
)
|
||||
self.addAsyncCleanup(client.close)
|
||||
|
||||
self.assertEqual((await client.keys())[0]["name"], "home")
|
||||
await client.create_key("living room")
|
||||
self.assertEqual(requests[1].url.raw_path, b"/key/living%20room")
|
||||
|
||||
async def test_reports_api_error_detail(self) -> None:
|
||||
def handler(_: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(409, json={"detail": "already exists"})
|
||||
|
||||
client = ManagementClient(
|
||||
"http://rex.test", transport=httpx.MockTransport(handler)
|
||||
)
|
||||
self.addAsyncCleanup(client.close)
|
||||
|
||||
with self.assertRaisesRegex(ManagementError, "already exists"):
|
||||
await client.create_device("desk")
|
||||
|
||||
|
||||
class DeviceManagerTests(unittest.TestCase):
|
||||
def test_action_changes_are_persisted(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "devices.json"
|
||||
manager = DeviceManager(path)
|
||||
manager.create_device("desk")
|
||||
manager.create_action("desk", "wake")
|
||||
|
||||
reloaded = DeviceManager(path)
|
||||
self.assertIsNotNone(reloaded.get_action("desk", "wake"))
|
||||
|
||||
reloaded.delete_action("desk", "wake")
|
||||
self.assertIsNone(DeviceManager(path).get_action("desk", "wake"))
|
||||
|
||||
|
||||
class FakeManagementClient:
|
||||
def __init__(self) -> None:
|
||||
self.created_keys: list[str] = []
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
async def keys(self) -> list[dict[str, str]]:
|
||||
return [{"name": "home", "secret": "abc"}]
|
||||
|
||||
async def devices(self) -> list[dict[str, object]]:
|
||||
return [{"name": "desk", "actions": [{"name": "wake"}]}]
|
||||
|
||||
async def create_key(self, name: str) -> None:
|
||||
self.created_keys.append(name)
|
||||
|
||||
|
||||
class ManagerAppTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_refresh_populates_keys_and_devices(self) -> None:
|
||||
app = RexManager("http://rex.test")
|
||||
app.client = FakeManagementClient() # type: ignore[assignment]
|
||||
|
||||
async with app.run_test() as pilot:
|
||||
await pilot.pause()
|
||||
keys = app.query_one("#keys", DataTable)
|
||||
devices = app.query_one("#devices", Tree)
|
||||
|
||||
self.assertEqual(keys.get_row_at(0), ["home", "abc"])
|
||||
self.assertEqual(devices.root.children[0].label.plain, "desk")
|
||||
self.assertEqual(devices.root.children[0].children[0].label.plain, "wake")
|
||||
|
||||
async def test_vim_navigation_and_create_prompt(self) -> None:
|
||||
app = RexManager("http://rex.test")
|
||||
client = FakeManagementClient()
|
||||
app.client = client # type: ignore[assignment]
|
||||
|
||||
async with app.run_test() as pilot:
|
||||
await pilot.press("n")
|
||||
await pilot.press(*"office", "enter")
|
||||
await pilot.pause()
|
||||
self.assertEqual(client.created_keys, ["office"])
|
||||
|
||||
await pilot.press("shift+l", "j")
|
||||
self.assertEqual(app.query_one(TabbedContent).active, "devices-tab")
|
||||
self.assertEqual(
|
||||
app.query_one("#devices", Tree).cursor_node.data,
|
||||
("device", "desk"),
|
||||
)
|
||||
|
||||
async def test_moves_between_key_columns_and_yanks_value(self) -> None:
|
||||
app = RexManager("http://rex.test")
|
||||
app.client = FakeManagementClient() # type: ignore[assignment]
|
||||
|
||||
async with app.run_test() as pilot:
|
||||
with patch.object(app, "copy_to_clipboard") as copy:
|
||||
await pilot.press("l", "y")
|
||||
|
||||
self.assertEqual(app.query_one("#keys", DataTable).cursor_column, 1)
|
||||
copy.assert_called_once_with("abc")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user