182 lines
6.9 KiB
Python
182 lines
6.9 KiB
Python
from __future__ import annotations
|
||
|
||
import os
|
||
from pathlib import Path
|
||
import pty
|
||
import select
|
||
import signal
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
|
||
from herdr_hub.remotes import RemoteEntry, RemoteQueryResult, RemoteRegistry, RemoteSession
|
||
from herdr_hub.selector import (
|
||
CatalogWindow,
|
||
PickerRow,
|
||
SessionSelector,
|
||
SelectorResult,
|
||
build_main_rows,
|
||
filter_rows,
|
||
)
|
||
|
||
|
||
class SelectorCatalogTests(unittest.TestCase):
|
||
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",
|
||
),
|
||
CatalogWindow(
|
||
"@2",
|
||
"remote:orangepi:work",
|
||
"R orangepi/work",
|
||
True,
|
||
False,
|
||
kind="remote",
|
||
target="orangepi",
|
||
session="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.assertIn("new:orangepi", by_key)
|
||
self.assertEqual(by_key["local:draft"].main, "○ draft")
|
||
self.assertEqual(by_key["close-active"].window_id, "@1")
|
||
|
||
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"])
|
||
|
||
|
||
class SelectorActionTests(unittest.TestCase):
|
||
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.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()
|