122 lines
4.8 KiB
Python
122 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import stat
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
|
|
from herdr_hub.remotes import (
|
|
RegistryError,
|
|
RemoteEntry,
|
|
RemoteQueryPool,
|
|
RemoteRegistry,
|
|
RemoteTargetError,
|
|
parse_remote_session_result,
|
|
remote_session_command,
|
|
ssh_destination_arguments,
|
|
)
|
|
|
|
|
|
class RemoteRegistryTests(unittest.TestCase):
|
|
def test_round_trip_update_remove_and_permissions(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "config" / "remotes.json"
|
|
registry = RemoteRegistry(path)
|
|
self.assertEqual(registry.load(), [])
|
|
self.assertTrue(registry.upsert(RemoteEntry("orangepi")))
|
|
self.assertFalse(registry.upsert(RemoteEntry("orangepi")))
|
|
self.assertTrue(registry.upsert(RemoteEntry("prod", "server")))
|
|
self.assertEqual(
|
|
registry.load(),
|
|
[RemoteEntry("orangepi"), RemoteEntry("prod", "server")],
|
|
)
|
|
self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600)
|
|
self.assertTrue(registry.remove("orangepi"))
|
|
self.assertFalse(registry.remove("missing"))
|
|
self.assertEqual(registry.load(), [RemoteEntry("prod", "server")])
|
|
|
|
def test_corrupt_registry_is_not_overwritten(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "remotes.json"
|
|
path.write_text("not json", encoding="utf-8")
|
|
registry = RemoteRegistry(path)
|
|
with self.assertRaises(RegistryError):
|
|
registry.upsert(RemoteEntry("orangepi"))
|
|
self.assertEqual(path.read_text(encoding="utf-8"), "not json")
|
|
|
|
|
|
class RemoteSshTests(unittest.TestCase):
|
|
def test_alias_and_url_targets_become_safe_ssh_arguments(self) -> None:
|
|
self.assertEqual(ssh_destination_arguments("orangepi"), ["orangepi"])
|
|
self.assertEqual(
|
|
ssh_destination_arguments("ssh://me@example.com:2222"),
|
|
["-p", "2222", "-l", "me", "example.com"],
|
|
)
|
|
command = remote_session_command("/usr/bin/ssh", RemoteEntry("host name"))
|
|
self.assertEqual(command[0], "/usr/bin/ssh")
|
|
self.assertIn("host name", command)
|
|
self.assertNotIn("sh", command[:1])
|
|
with self.assertRaises(RemoteTargetError):
|
|
ssh_destination_arguments("ssh://me:secret@example.com")
|
|
with self.assertRaises(RemoteTargetError):
|
|
ssh_destination_arguments("ssh://example.com/a/path")
|
|
|
|
def test_remote_session_json_keeps_running_and_stopped(self) -> None:
|
|
result = parse_remote_session_result(
|
|
RemoteEntry("orangepi"),
|
|
0,
|
|
'{"sessions":['
|
|
'{"name":"old","running":false},'
|
|
'{"name":"default","running":true}]}',
|
|
"",
|
|
)
|
|
self.assertIsNone(result.error)
|
|
self.assertEqual(
|
|
[(session.name, session.running) for session in result.sessions],
|
|
[("default", True), ("old", False)],
|
|
)
|
|
|
|
def test_query_pool_reports_results_and_spawn_failures(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
fake_ssh = Path(directory) / "ssh"
|
|
fake_ssh.write_text(
|
|
"#!/bin/sh\n"
|
|
"printf '%s\\n' '{\"sessions\":[{\"name\":\"work\",\"running\":true}]}'\n",
|
|
encoding="utf-8",
|
|
)
|
|
fake_ssh.chmod(0o755)
|
|
pool = RemoteQueryPool(str(fake_ssh), query_timeout=1)
|
|
self.addCleanup(pool.close)
|
|
pool.enqueue(RemoteEntry("orangepi"))
|
|
results = []
|
|
deadline = time.monotonic() + 2
|
|
while not results and time.monotonic() < deadline:
|
|
results = pool.poll()
|
|
time.sleep(0.01)
|
|
self.assertEqual(results[0].sessions[0].name, "work")
|
|
|
|
missing = RemoteQueryPool(str(Path(directory) / "missing"))
|
|
self.addCleanup(missing.close)
|
|
missing.enqueue(RemoteEntry("orangepi"))
|
|
failure = missing.poll()
|
|
self.assertEqual(len(failure), 1)
|
|
self.assertIn("cannot start", failure[0].error or "")
|
|
|
|
def test_query_pool_times_out_without_blocking_other_results(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
slow_ssh = Path(directory) / "ssh"
|
|
slow_ssh.write_text("#!/bin/sh\nexec sleep 5\n", encoding="utf-8")
|
|
slow_ssh.chmod(0o755)
|
|
pool = RemoteQueryPool(str(slow_ssh), query_timeout=0.05)
|
|
self.addCleanup(pool.close)
|
|
pool.enqueue(RemoteEntry("slow"))
|
|
time.sleep(0.08)
|
|
results = pool.poll()
|
|
self.assertEqual(len(results), 1)
|
|
self.assertIn("timed out", results[0].error or "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|