Improve remote session management and Windows support
This commit is contained in:
+181
-3
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
import stat
|
||||
import tempfile
|
||||
@@ -7,14 +8,21 @@ import time
|
||||
import unittest
|
||||
|
||||
from herdr_hub.remotes import (
|
||||
REMOTE_DIRECT_COMMAND,
|
||||
REMOTE_LIST_SCRIPT,
|
||||
RegistryError,
|
||||
RemoteEntry,
|
||||
RemoteQueryPool,
|
||||
RemoteRegistry,
|
||||
RemoteTargetError,
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@@ -26,15 +34,31 @@ class RemoteRegistryTests(unittest.TestCase):
|
||||
self.assertEqual(registry.load(), [])
|
||||
self.assertTrue(registry.upsert(RemoteEntry("orangepi")))
|
||||
self.assertFalse(registry.upsert(RemoteEntry("orangepi")))
|
||||
self.assertTrue(registry.upsert(RemoteEntry("prod", "server")))
|
||||
self.assertTrue(registry.upsert(RemoteEntry("prod", "server", "ssh")))
|
||||
self.assertEqual(
|
||||
registry.load(),
|
||||
[RemoteEntry("orangepi"), RemoteEntry("prod", "server")],
|
||||
[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")])
|
||||
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:
|
||||
@@ -45,6 +69,25 @@ class RemoteRegistryTests(unittest.TestCase):
|
||||
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:
|
||||
@@ -56,12 +99,93 @@ class RemoteSshTests(unittest.TestCase):
|
||||
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"),
|
||||
@@ -77,6 +201,15 @@ class RemoteSshTests(unittest.TestCase):
|
||||
[("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"
|
||||
@@ -103,6 +236,51 @@ class RemoteSshTests(unittest.TestCase):
|
||||
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_query_pool_times_out_without_blocking_other_results(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
slow_ssh = Path(directory) / "ssh"
|
||||
|
||||
Reference in New Issue
Block a user