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()