Files
herdr-session-switch/tests/test_remotes.py
T
gitea-adminandClaude Fable 5.1 3f11734ce4 Decode streamed update output at character boundaries
The update popup kept the last 64 KiB of raw bytes and decoded the
whole buffer on every refresh. Cutting the buffer, or reading a chunk
that ended inside a multi-byte sequence, made the UTF-8 decode fail and
sent everything through the GB18030 fallback, garbling the display.

Clip the retained tail at a UTF-8 boundary and decode still-growing
streams with an incremental decoder that holds back an incomplete
trailing sequence instead of treating it as corruption.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-10-03 03:00:09 +00:00

325 lines
13 KiB
Python
Raw 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_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()