Files
gitea-adminandClaude Fable 5.1 72d594033e Stop cancelled remote queries without blocking the selector
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>
2026-10-03 03:04:27 +00:00

352 lines
14 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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()