RemoteQueryPool.cancel() terminated the SSH process and then waited up to 250 ms for it, on the curses thread. Ctrl+R re-enqueues every host, so a refresh of six hosts could freeze the popup for over a second. Only send SIGTERM when cancelling; reap the process on later polls, escalating to SIGKILL after 250 ms and closing its pipes once it has exited. close() still waits, since the popup is exiting anyway. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
352 lines
14 KiB
Python
352 lines
14 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
from pathlib import Path
|
||
import stat
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
|
||
from herdr_hub.remotes import (
|
||
REMOTE_DIRECT_COMMAND,
|
||
REMOTE_LIST_SCRIPT,
|
||
RegistryError,
|
||
RemoteEntry,
|
||
RemoteQueryPool,
|
||
RemoteRegistry,
|
||
RemoteTargetError,
|
||
clip_utf8_tail,
|
||
delete_remote_session,
|
||
decode_remote_output,
|
||
parse_remote_session_result,
|
||
remote_delete_session_command,
|
||
remote_session_command,
|
||
remote_stop_session_command,
|
||
ssh_destination_arguments,
|
||
stop_remote_session,
|
||
)
|
||
|
||
|
||
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", "ssh")))
|
||
self.assertEqual(
|
||
registry.load(),
|
||
[RemoteEntry("orangepi"), RemoteEntry("prod", "server", "ssh")],
|
||
)
|
||
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", "ssh")],
|
||
)
|
||
|
||
def test_discovery_does_not_overwrite_preferences_and_ssh_is_sticky(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
||
registry.upsert(RemoteEntry("host", "server"))
|
||
registry.remember_many([RemoteEntry("host", "local")])
|
||
self.assertEqual(registry.load(), [RemoteEntry("host", "server")])
|
||
|
||
registry.remember_many([RemoteEntry("host", "local", "ssh")])
|
||
self.assertEqual(
|
||
registry.load(),
|
||
[RemoteEntry("host", "server", "ssh")],
|
||
)
|
||
|
||
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")
|
||
|
||
def test_legacy_registry_defaults_to_native_transport(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
path = Path(directory) / "remotes.json"
|
||
path.write_text(
|
||
json.dumps(
|
||
{
|
||
"version": 1,
|
||
"remotes": [
|
||
{"target": "legacy", "remote_keybindings": "server"}
|
||
],
|
||
}
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
self.assertEqual(
|
||
RemoteRegistry(path).load(),
|
||
[RemoteEntry("legacy", "server", "native")],
|
||
)
|
||
|
||
|
||
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.assertEqual(command[-1], REMOTE_DIRECT_COMMAND)
|
||
fallback = remote_session_command(
|
||
"/usr/bin/ssh",
|
||
RemoteEntry("orangepi"),
|
||
posix_fallback=True,
|
||
)
|
||
self.assertEqual(fallback[-1], REMOTE_LIST_SCRIPT)
|
||
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_delete_stopped_remote_session_uses_safe_command_and_posix_fallback(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
fake_ssh = Path(directory) / "ssh"
|
||
fake_ssh.write_text(
|
||
"#!/bin/sh\n"
|
||
"last=''\n"
|
||
"for argument do last=$argument; done\n"
|
||
"case \"$last\" in\n"
|
||
" *\"if [ -x\"*)\n"
|
||
" printf '%s\\n' '{\"deleted\":true}'\n"
|
||
" exit 0\n"
|
||
" ;;\n"
|
||
"esac\n"
|
||
"exit 127\n",
|
||
encoding="utf-8",
|
||
)
|
||
fake_ssh.chmod(0o755)
|
||
delete_remote_session(
|
||
str(fake_ssh),
|
||
RemoteEntry("orangepi"),
|
||
"old-session",
|
||
timeout=1,
|
||
)
|
||
|
||
command = remote_delete_session_command(
|
||
"/usr/bin/ssh",
|
||
RemoteEntry("orangepi"),
|
||
"old-session",
|
||
)
|
||
self.assertEqual(
|
||
command[-1],
|
||
"herdr session delete --json -- old-session",
|
||
)
|
||
with self.assertRaises(ValueError):
|
||
remote_delete_session_command(
|
||
"/usr/bin/ssh",
|
||
RemoteEntry("orangepi"),
|
||
"bad session",
|
||
)
|
||
|
||
def test_stop_running_remote_session_uses_safe_command_and_posix_fallback(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
fake_ssh = Path(directory) / "ssh"
|
||
fake_ssh.write_text(
|
||
"#!/bin/sh\n"
|
||
"last=''\n"
|
||
"for argument do last=$argument; done\n"
|
||
"case \"$last\" in\n"
|
||
" *\"if [ -x\"*) exit 0 ;;\n"
|
||
"esac\n"
|
||
"exit 127\n",
|
||
encoding="utf-8",
|
||
)
|
||
fake_ssh.chmod(0o755)
|
||
stop_remote_session(
|
||
str(fake_ssh),
|
||
RemoteEntry("orangepi"),
|
||
"work",
|
||
timeout=1,
|
||
)
|
||
|
||
command = remote_stop_session_command(
|
||
"/usr/bin/ssh",
|
||
RemoteEntry("orangepi"),
|
||
"work",
|
||
)
|
||
self.assertEqual(command[-1], "herdr session stop -- work")
|
||
with self.assertRaises(ValueError):
|
||
remote_stop_session_command(
|
||
"/usr/bin/ssh",
|
||
RemoteEntry("orangepi"),
|
||
"bad session",
|
||
)
|
||
|
||
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)],
|
||
)
|
||
|
||
windows = parse_remote_session_result(
|
||
RemoteEntry("windows"),
|
||
0,
|
||
'{"sessions":[{"name":"default","running":true,'
|
||
'"session_dir":"C:\\\\Users\\\\me\\\\AppData\\\\Roaming\\\\herdr"}]}',
|
||
"",
|
||
)
|
||
self.assertEqual(windows.platform, "windows")
|
||
|
||
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_falls_back_for_posix_path_and_decodes_windows_errors(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
root = Path(directory)
|
||
fallback_ssh = root / "fallback-ssh"
|
||
fallback_ssh.write_text(
|
||
"#!/bin/sh\n"
|
||
"last=''\n"
|
||
"for argument do last=$argument; done\n"
|
||
"case \"$last\" in\n"
|
||
" *\"if [ -x\"*)\n"
|
||
" printf '%s\\n' '{\"sessions\":[{\"name\":\"work\",\"running\":true}]}'\n"
|
||
" exit 0\n"
|
||
" ;;\n"
|
||
"esac\n"
|
||
"exit 127\n",
|
||
encoding="utf-8",
|
||
)
|
||
fallback_ssh.chmod(0o755)
|
||
pool = RemoteQueryPool(str(fallback_ssh), query_timeout=1)
|
||
self.addCleanup(pool.close)
|
||
pool.enqueue(RemoteEntry("linux-host"))
|
||
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")
|
||
|
||
windows_ssh = root / "windows-ssh"
|
||
windows_ssh.write_bytes(
|
||
b"#!/bin/sh\nprintf '\\264\\355\\316\\363' >&2\nexit 1\n"
|
||
)
|
||
windows_ssh.chmod(0o755)
|
||
windows_pool = RemoteQueryPool(str(windows_ssh), query_timeout=1)
|
||
self.addCleanup(windows_pool.close)
|
||
windows_pool.enqueue(RemoteEntry("windows-host"))
|
||
windows_results = []
|
||
deadline = time.monotonic() + 2
|
||
while not windows_results and time.monotonic() < deadline:
|
||
windows_results = windows_pool.poll()
|
||
time.sleep(0.01)
|
||
self.assertIn("错误", windows_results[0].error or "")
|
||
|
||
self.assertEqual(decode_remote_output("错误".encode("gbk")), "错误")
|
||
|
||
def test_cancel_does_not_block_and_reaps_the_process_later(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
slow_ssh = Path(directory) / "ssh"
|
||
# Ignore SIGTERM so the pool has to escalate to SIGKILL.
|
||
slow_ssh.write_text("#!/bin/sh\ntrap '' TERM\nsleep 5 & wait\n", encoding="utf-8")
|
||
slow_ssh.chmod(0o755)
|
||
pool = RemoteQueryPool(str(slow_ssh), query_timeout=5)
|
||
self.addCleanup(pool.close)
|
||
pool.enqueue(RemoteEntry("slow"))
|
||
time.sleep(0.05)
|
||
process = pool._active["slow"].process # type: ignore[attr-defined]
|
||
|
||
started = time.monotonic()
|
||
pool.cancel("slow")
|
||
self.assertLess(time.monotonic() - started, 0.1)
|
||
self.assertFalse(pool.is_loading("slow"))
|
||
|
||
deadline = time.monotonic() + 3
|
||
while process.poll() is None and time.monotonic() < deadline:
|
||
pool.poll()
|
||
time.sleep(0.02)
|
||
self.assertIsNotNone(process.poll())
|
||
pool.poll()
|
||
self.assertEqual(pool._stopping, []) # type: ignore[attr-defined]
|
||
assert process.stdout is not None
|
||
self.assertTrue(process.stdout.closed)
|
||
|
||
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 "")
|
||
|
||
|
||
class RemoteOutputDecodingTests(unittest.TestCase):
|
||
def test_partial_streams_hold_back_incomplete_sequences(self) -> None:
|
||
text = "更新目标 Session:default\n正在热切换…"
|
||
raw = text.encode("utf-8")
|
||
# A chunk boundary inside the last character must not garble the
|
||
# characters already received.
|
||
self.assertEqual(decode_remote_output(raw[:-1], partial=True), text[:-1])
|
||
self.assertEqual(decode_remote_output(raw, partial=True), text)
|
||
gbk = "错误:下载失败".encode("gb18030")
|
||
self.assertEqual(decode_remote_output(gbk[:-1], partial=True), "错误:下载失")
|
||
# Final decoding keeps its strict behaviour.
|
||
self.assertEqual(decode_remote_output(raw), text)
|
||
self.assertEqual(decode_remote_output(gbk), "错误:下载失败")
|
||
|
||
def test_clip_utf8_tail_never_starts_inside_a_character(self) -> None:
|
||
raw = "abc更新目标".encode("utf-8")
|
||
self.assertEqual(clip_utf8_tail(raw, len(raw)), raw)
|
||
# Cutting one byte into "更" drops the rest of that character.
|
||
clipped = clip_utf8_tail(raw, len(raw) - 4)
|
||
self.assertEqual(clipped.decode("utf-8"), "新目标")
|
||
self.assertEqual(clip_utf8_tail(raw, 3).decode("utf-8"), "标")
|
||
self.assertEqual(clip_utf8_tail(b"", 10), b"")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|