563 lines
22 KiB
Python
563 lines
22 KiB
Python
from __future__ import annotations
|
||
|
||
import curses
|
||
from dataclasses import replace
|
||
import os
|
||
from pathlib import Path
|
||
import pty
|
||
import select
|
||
import signal
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from unittest.mock import patch
|
||
|
||
from herdr_hub.remotes import RemoteEntry, RemoteQueryResult, RemoteRegistry, RemoteSession
|
||
from herdr_hub.selector import (
|
||
CatalogWindow,
|
||
DeleteSessionOutcome,
|
||
PickerRow,
|
||
SessionSelector,
|
||
SelectorResult,
|
||
build_main_rows,
|
||
filter_rows,
|
||
input_display_line,
|
||
pad_terminal_text,
|
||
terminal_display_width,
|
||
truncate_terminal_text,
|
||
)
|
||
|
||
|
||
def wait_for_action(selector: SessionSelector) -> None:
|
||
deadline = time.monotonic() + 2
|
||
while selector.action_future is not None and time.monotonic() < deadline:
|
||
selector._poll_action() # type: ignore[attr-defined]
|
||
time.sleep(0.01)
|
||
if selector.action_future is not None:
|
||
raise AssertionError("selector action did not finish")
|
||
|
||
|
||
class SelectorCatalogTests(unittest.TestCase):
|
||
def test_terminal_display_width_counts_cjk_and_combining_characters(self) -> None:
|
||
self.assertEqual(terminal_display_width("Rename: work"), 12)
|
||
self.assertEqual(terminal_display_width("显示名称: 工作"), 14)
|
||
self.assertEqual(terminal_display_width("e\N{COMBINING ACUTE ACCENT}"), 1)
|
||
self.assertEqual(truncate_terminal_text("主机abc", 5), "主机a")
|
||
self.assertEqual(terminal_display_width(pad_terminal_text("主机", 6)), 6)
|
||
self.assertEqual(input_display_line("名称", "abcdefgh", 9), "名称: …gh")
|
||
|
||
def test_rows_merge_open_running_stopped_and_new_sessions(self) -> None:
|
||
windows = [
|
||
CatalogWindow(
|
||
"@1",
|
||
"local:default",
|
||
"L default",
|
||
False,
|
||
True,
|
||
kind="local",
|
||
session="default",
|
||
default_label="L default",
|
||
),
|
||
CatalogWindow(
|
||
"@2",
|
||
"remote:orangepi:work",
|
||
"R orangepi/work",
|
||
True,
|
||
False,
|
||
kind="remote",
|
||
target="orangepi",
|
||
session="work",
|
||
default_label="R orangepi/work",
|
||
),
|
||
]
|
||
result = RemoteQueryResult(
|
||
RemoteEntry("orangepi"),
|
||
sessions=(
|
||
RemoteSession("default", True),
|
||
RemoteSession("old", False),
|
||
RemoteSession("work", True),
|
||
),
|
||
)
|
||
rows = build_main_rows(
|
||
windows,
|
||
[RemoteEntry("orangepi")],
|
||
{"orangepi": result},
|
||
set(),
|
||
local_sessions=["default", "draft"],
|
||
)
|
||
by_key = {row.key: row for row in rows}
|
||
self.assertEqual(by_key["window:@1"].main, "▶ L default")
|
||
self.assertEqual(by_key["window:@2"].main, "× R orangepi/work")
|
||
self.assertIn("remote:orangepi:default", by_key)
|
||
self.assertIn("remote:orangepi:old", by_key)
|
||
self.assertTrue(by_key["remote:orangepi:old"].deletable)
|
||
self.assertIn("d/D 删除", by_key["remote:orangepi:old"].detail)
|
||
self.assertFalse(by_key["remote:orangepi:default"].deletable)
|
||
self.assertIn("new:orangepi", by_key)
|
||
self.assertEqual(by_key["local:draft"].main, "○ draft")
|
||
self.assertEqual(by_key["close-active"].window_id, "@1")
|
||
self.assertIn("manage-sessions", by_key)
|
||
|
||
def test_loading_error_and_filter_rows(self) -> None:
|
||
entry = RemoteEntry("prod")
|
||
loading = build_main_rows([], [entry], {}, {"prod"})
|
||
self.assertTrue(any(row.kind == "loading" for row in loading))
|
||
failed = build_main_rows(
|
||
[],
|
||
[entry],
|
||
{"prod": RemoteQueryResult(entry, error="offline")},
|
||
set(),
|
||
)
|
||
self.assertTrue(any(row.kind == "retry" for row in failed))
|
||
matches = filter_rows(failed, "offline")
|
||
self.assertEqual([row.kind for row in matches], ["retry"])
|
||
|
||
def test_active_remote_session_has_close_and_stop_action(self) -> None:
|
||
local = CatalogWindow(
|
||
"@1",
|
||
"local:default",
|
||
"L default",
|
||
False,
|
||
False,
|
||
kind="local",
|
||
session="default",
|
||
)
|
||
remote = CatalogWindow(
|
||
"@2",
|
||
"remote:orangepi:work",
|
||
"R orangepi/work",
|
||
False,
|
||
True,
|
||
kind="remote",
|
||
target="orangepi",
|
||
session="work",
|
||
)
|
||
rows = build_main_rows(
|
||
[local, remote],
|
||
[RemoteEntry("orangepi")],
|
||
{},
|
||
{"orangepi"},
|
||
)
|
||
stop = next(row for row in rows if row.kind == "stop-session")
|
||
self.assertTrue(stop.selectable)
|
||
self.assertEqual(stop.window_id, "@2")
|
||
self.assertIn("任务会中断", stop.detail)
|
||
|
||
last_tab_rows = build_main_rows(
|
||
[remote],
|
||
[RemoteEntry("orangepi")],
|
||
{},
|
||
{"orangepi"},
|
||
)
|
||
last_stop = next(
|
||
row for row in last_tab_rows if row.kind == "stop-session"
|
||
)
|
||
self.assertTrue(last_stop.selectable)
|
||
self.assertIn("Tab 将保留", last_stop.detail)
|
||
|
||
def test_windows_remote_migrates_to_ssh_transport(self) -> None:
|
||
window = CatalogWindow(
|
||
"@4",
|
||
"remote:windows:default",
|
||
"R windows/default",
|
||
True,
|
||
False,
|
||
kind="remote",
|
||
target="windows",
|
||
session="default",
|
||
remote_transport="native",
|
||
)
|
||
result = RemoteQueryResult(
|
||
RemoteEntry("windows"),
|
||
sessions=(RemoteSession("default", True),),
|
||
platform="windows",
|
||
)
|
||
rows = build_main_rows(
|
||
[window],
|
||
[RemoteEntry("windows")],
|
||
{"windows": result},
|
||
set(),
|
||
)
|
||
remote = next(row for row in rows if row.kind == "remote")
|
||
self.assertEqual(remote.remote_transport, "ssh")
|
||
self.assertIn("切换为 SSH 模式", remote.detail)
|
||
|
||
def test_existing_ssh_window_is_never_downgraded_while_query_is_unknown(self) -> None:
|
||
entry = RemoteEntry("windows")
|
||
window = CatalogWindow(
|
||
"@4",
|
||
"remote:windows:default",
|
||
"windows",
|
||
False,
|
||
True,
|
||
kind="remote",
|
||
target="windows",
|
||
session="default",
|
||
remote_transport="ssh",
|
||
)
|
||
for results, loading in [
|
||
({}, {"windows"}),
|
||
({"windows": RemoteQueryResult(entry, error="offline")}, set()),
|
||
({"windows": RemoteQueryResult(entry)}, set()),
|
||
]:
|
||
rows = build_main_rows([window], [entry], results, loading)
|
||
row = next(item for item in rows if item.key == "window:@4")
|
||
self.assertEqual(row.kind, "window")
|
||
|
||
def test_explicit_ssh_entry_supports_an_empty_windows_host(self) -> None:
|
||
entry = RemoteEntry("windows", remote_transport="ssh")
|
||
rows = build_main_rows(
|
||
[],
|
||
[entry],
|
||
{"windows": RemoteQueryResult(entry)},
|
||
set(),
|
||
)
|
||
new_row = next(item for item in rows if item.kind == "new")
|
||
self.assertEqual(new_row.remote_transport, "ssh")
|
||
|
||
def test_windows_new_session_keeps_the_row_transport(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
||
entry = RemoteEntry("windows")
|
||
registry.upsert(entry)
|
||
selector = SessionSelector([], registry, "/missing/ssh")
|
||
self.addCleanup(selector.pool.close)
|
||
selector.results["windows"] = RemoteQueryResult(
|
||
entry,
|
||
sessions=(RemoteSession("default", True),),
|
||
platform="windows",
|
||
)
|
||
new_row = next(
|
||
item
|
||
for item in selector._rows() # type: ignore[attr-defined]
|
||
if item.kind == "new"
|
||
)
|
||
selector._activate(new_row) # type: ignore[attr-defined]
|
||
result = selector._finish_new("work") # type: ignore[attr-defined]
|
||
self.assertEqual(
|
||
result,
|
||
SelectorResult.remote("windows", "work", "local", "ssh"),
|
||
)
|
||
|
||
|
||
class SelectorActionTests(unittest.TestCase):
|
||
def test_remote_session_can_be_stopped_and_closed_without_closing_popup(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
||
registry.upsert(RemoteEntry("orangepi"))
|
||
local = CatalogWindow(
|
||
"@1",
|
||
"local:default",
|
||
"L default",
|
||
False,
|
||
False,
|
||
kind="local",
|
||
session="default",
|
||
)
|
||
remote = CatalogWindow(
|
||
"@2",
|
||
"remote:orangepi:work",
|
||
"R orangepi/work",
|
||
False,
|
||
True,
|
||
kind="remote",
|
||
target="orangepi",
|
||
session="work",
|
||
)
|
||
stopped: list[tuple[str, str, str]] = []
|
||
|
||
def stop_handler(
|
||
window_id: str,
|
||
target: str,
|
||
session: str,
|
||
) -> list[CatalogWindow]:
|
||
stopped.append((window_id, target, session))
|
||
return [replace(local, active=True)]
|
||
|
||
selector = SessionSelector(
|
||
[local, remote],
|
||
registry,
|
||
str(Path(directory) / "missing-ssh"),
|
||
stop_session_handler=stop_handler,
|
||
)
|
||
self.addCleanup(selector.pool.close)
|
||
self.addCleanup(selector.action_executor.shutdown, wait=True)
|
||
selector.results["orangepi"] = RemoteQueryResult(
|
||
RemoteEntry("orangepi"),
|
||
sessions=(RemoteSession("work", True),),
|
||
)
|
||
|
||
stop_row = next(
|
||
row
|
||
for row in selector._rows() # type: ignore[attr-defined]
|
||
if row.kind == "stop-session"
|
||
)
|
||
selector._activate(stop_row) # type: ignore[attr-defined]
|
||
self.assertEqual(selector.view, "confirm-stop-session")
|
||
confirm = next(
|
||
row
|
||
for row in selector._rows() # type: ignore[attr-defined]
|
||
if row.kind == "stop-session-yes"
|
||
)
|
||
selector._activate(confirm) # type: ignore[attr-defined]
|
||
self.assertEqual(selector.view, "working-stop-session")
|
||
wait_for_action(selector)
|
||
|
||
self.assertEqual(stopped, [("@2", "orangepi", "work")])
|
||
self.assertEqual(selector.view, "main")
|
||
self.assertEqual([window.window_id for window in selector.windows], ["@1"])
|
||
self.assertFalse(selector.results["orangepi"].sessions[0].running)
|
||
self.assertIn("已关闭 Tab 并停止", selector.message)
|
||
|
||
def test_stopped_remote_session_can_be_deleted_without_closing_popup(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
||
registry.upsert(RemoteEntry("orangepi"))
|
||
deleted: list[tuple[str, str]] = []
|
||
selector = SessionSelector(
|
||
[],
|
||
registry,
|
||
str(Path(directory) / "missing-ssh"),
|
||
delete_session_handler=lambda target, session: deleted.append(
|
||
(target, session)
|
||
),
|
||
)
|
||
self.addCleanup(selector.pool.close)
|
||
self.addCleanup(selector.action_executor.shutdown, wait=True)
|
||
selector.results["orangepi"] = RemoteQueryResult(
|
||
RemoteEntry("orangepi"),
|
||
sessions=(
|
||
RemoteSession("default", True),
|
||
RemoteSession("old", False),
|
||
),
|
||
)
|
||
|
||
rows = selector._rows() # type: ignore[attr-defined]
|
||
running = next(
|
||
row for row in rows if row.key == "remote:orangepi:default"
|
||
)
|
||
selector.selected_key = running.key
|
||
selector._handle_key("d", rows) # type: ignore[attr-defined]
|
||
self.assertEqual(selector.view, "main")
|
||
self.assertEqual(deleted, [])
|
||
self.assertEqual(selector.query, "")
|
||
self.assertIn("只能删除", selector.message)
|
||
|
||
selector.query = ""
|
||
old = next(row for row in rows if row.key == "remote:orangepi:old")
|
||
selector.selected_key = old.key
|
||
selector._handle_key("D", rows) # type: ignore[attr-defined]
|
||
self.assertEqual(selector.view, "confirm-delete-session")
|
||
confirm = next(
|
||
row
|
||
for row in selector._rows() # type: ignore[attr-defined]
|
||
if row.kind == "delete-session-yes"
|
||
)
|
||
selector._drawn_rows[6] = confirm # type: ignore[attr-defined]
|
||
with patch(
|
||
"herdr_hub.selector.curses.getmouse",
|
||
return_value=(0, 0, 6, 0, curses.BUTTON1_RELEASED),
|
||
):
|
||
selector._handle_mouse( # type: ignore[attr-defined]
|
||
selector._rows() # type: ignore[attr-defined]
|
||
)
|
||
self.assertEqual(selector.view, "working-delete-session")
|
||
wait_for_action(selector)
|
||
|
||
self.assertEqual(deleted, [("orangepi", "old")])
|
||
self.assertEqual(selector.view, "main")
|
||
self.assertIn("已删除历史 Session", selector.message)
|
||
self.assertNotIn(
|
||
"old",
|
||
[item.name for item in selector.results["orangepi"].sessions],
|
||
)
|
||
|
||
def test_delete_partial_success_is_reported_as_warning_not_failure(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
||
registry.upsert(RemoteEntry("host"))
|
||
selector = SessionSelector(
|
||
[],
|
||
registry,
|
||
"/missing/ssh",
|
||
delete_session_handler=lambda _target, _session: DeleteSessionOutcome(
|
||
"显示名称清理失败"
|
||
),
|
||
)
|
||
self.addCleanup(selector.pool.close)
|
||
self.addCleanup(selector.action_executor.shutdown, wait=True)
|
||
selector.results["host"] = RemoteQueryResult(
|
||
RemoteEntry("host"),
|
||
sessions=(RemoteSession("old", False),),
|
||
)
|
||
selector._finish_delete_session("host", "old") # type: ignore[attr-defined]
|
||
wait_for_action(selector)
|
||
self.assertEqual(selector.view, "main")
|
||
self.assertIn("已删除历史 Session", selector.message)
|
||
self.assertIn("显示名称清理失败", selector.message)
|
||
|
||
def test_dead_window_reconnect_action_and_new_session_action(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
registry = RemoteRegistry(Path(directory) / "remotes.json")
|
||
registry.upsert(RemoteEntry("orangepi", "server"))
|
||
selector = SessionSelector(
|
||
[CatalogWindow("@2", "key", "remote", True, False)],
|
||
registry,
|
||
str(Path(directory) / "missing-ssh"),
|
||
)
|
||
self.addCleanup(selector.pool.close)
|
||
window_result = selector._activate( # type: ignore[attr-defined]
|
||
next(row for row in selector._rows() if row.kind == "window") # type: ignore[attr-defined]
|
||
)
|
||
self.assertEqual(window_result, SelectorResult.window("@2", dead=True))
|
||
|
||
close_result = selector._activate( # type: ignore[attr-defined]
|
||
PickerRow("close", "close", "close", window_id="@2")
|
||
)
|
||
self.assertEqual(close_result, SelectorResult.close("@2"))
|
||
|
||
local_result = selector._activate( # type: ignore[attr-defined]
|
||
PickerRow("local", "local", "local", session="draft")
|
||
)
|
||
self.assertEqual(local_result, SelectorResult.local("draft"))
|
||
|
||
selector = SessionSelector(
|
||
[
|
||
CatalogWindow(
|
||
"@3",
|
||
"local:default",
|
||
"日常工作",
|
||
False,
|
||
True,
|
||
kind="local",
|
||
session="default",
|
||
default_label="L default",
|
||
)
|
||
],
|
||
registry,
|
||
str(Path(directory) / "missing-ssh"),
|
||
aliases={"local:default": "日常工作"},
|
||
)
|
||
self.addCleanup(selector.pool.close)
|
||
selector._activate( # type: ignore[attr-defined]
|
||
next(
|
||
row
|
||
for row in selector._rows() # type: ignore[attr-defined]
|
||
if row.kind == "manage-sessions"
|
||
)
|
||
)
|
||
rename_row = next(
|
||
row
|
||
for row in selector._rows() # type: ignore[attr-defined]
|
||
if row.kind == "rename"
|
||
)
|
||
selector._activate(rename_row) # type: ignore[attr-defined]
|
||
self.assertEqual(selector.view, "input-rename")
|
||
self.assertEqual(selector.input_value, "日常工作")
|
||
self.assertEqual(
|
||
selector._finish_rename(" 新名称 "), # type: ignore[attr-defined]
|
||
SelectorResult.rename("local:default", "新名称"),
|
||
)
|
||
self.assertEqual(
|
||
selector._finish_rename(""), # type: ignore[attr-defined]
|
||
SelectorResult.rename("local:default", ""),
|
||
)
|
||
|
||
renamed: list[tuple[str, str]] = []
|
||
live_selector = SessionSelector(
|
||
selector.windows,
|
||
registry,
|
||
str(Path(directory) / "missing-ssh"),
|
||
aliases={"local:default": "日常工作"},
|
||
rename_handler=lambda key, alias: renamed.append((key, alias)),
|
||
)
|
||
self.addCleanup(live_selector.pool.close)
|
||
live_selector._show_manage_sessions() # type: ignore[attr-defined]
|
||
live_selector._activate( # type: ignore[attr-defined]
|
||
next(
|
||
row
|
||
for row in live_selector._rows() # type: ignore[attr-defined]
|
||
if row.kind == "rename"
|
||
)
|
||
)
|
||
self.assertIsNone(
|
||
live_selector._finish_rename("连续改名") # type: ignore[attr-defined]
|
||
)
|
||
self.assertEqual(renamed, [("local:default", "连续改名")])
|
||
self.assertEqual(live_selector.view, "manage-sessions")
|
||
self.assertEqual(live_selector.selected_key, "rename:local:default")
|
||
self.assertEqual(live_selector.aliases["local:default"], "连续改名")
|
||
self.assertIn("已保存", live_selector.message)
|
||
|
||
selector.view = "input-new"
|
||
selector.input_target = "orangepi"
|
||
new_result = selector._finish_new("work") # type: ignore[attr-defined]
|
||
self.assertEqual(
|
||
new_result,
|
||
SelectorResult.remote("orangepi", "work", "server"),
|
||
)
|
||
self.assertIsNone(selector._finish_new("bad name")) # type: ignore[attr-defined]
|
||
|
||
@unittest.skipUnless(hasattr(pty, "fork"), "a Unix PTY is required")
|
||
def test_curses_selector_handles_single_click_and_esc(self) -> None:
|
||
with tempfile.TemporaryDirectory() as directory:
|
||
pid, descriptor = pty.fork()
|
||
if pid == 0:
|
||
os.environ["TERM"] = "xterm-256color"
|
||
selector = SessionSelector(
|
||
[],
|
||
RemoteRegistry(Path(directory) / "remotes.json"),
|
||
"/missing/ssh",
|
||
)
|
||
try:
|
||
result = selector.run()
|
||
os._exit(0 if result is None else 2)
|
||
except BaseException:
|
||
os._exit(3)
|
||
|
||
status: int | None = None
|
||
try:
|
||
deadline = time.monotonic() + 3
|
||
output = b""
|
||
while time.monotonic() < deadline and b"Sessions" not in output:
|
||
readable, _, _ = select.select([descriptor], [], [], 0.1)
|
||
if readable:
|
||
output += os.read(descriptor, 4096)
|
||
self.assertIn(b"Sessions", output)
|
||
# Empty registries render "+ 添加远端" as the first selectable
|
||
# row at terminal row 5. A single SGR mouse click must activate it.
|
||
os.write(descriptor, b"\x1b[<0;5;5M\x1b[<0;5;5m")
|
||
opened_input = b""
|
||
input_deadline = time.monotonic() + 2
|
||
while time.monotonic() < input_deadline and b"SSH target" not in opened_input:
|
||
readable, _, _ = select.select([descriptor], [], [], 0.1)
|
||
if readable:
|
||
opened_input += os.read(descriptor, 4096)
|
||
self.assertIn(b"SSH target", opened_input)
|
||
os.write(descriptor, b"\x1b\x1b")
|
||
while time.monotonic() < deadline:
|
||
completed, wait_status = os.waitpid(pid, os.WNOHANG)
|
||
if completed == pid:
|
||
status = wait_status
|
||
break
|
||
time.sleep(0.05)
|
||
self.assertIsNotNone(status, "selector did not exit after Esc")
|
||
assert status is not None
|
||
self.assertTrue(os.WIFEXITED(status))
|
||
self.assertEqual(os.WEXITSTATUS(status), 0)
|
||
finally:
|
||
try:
|
||
os.close(descriptor)
|
||
except OSError:
|
||
pass
|
||
if status is None:
|
||
try:
|
||
os.kill(pid, signal.SIGKILL)
|
||
except ProcessLookupError:
|
||
pass
|
||
os.waitpid(pid, 0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|