Files

636 lines
22 KiB
Python

from __future__ import annotations
import base64
import importlib.util
import json
import os
import subprocess
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
MODULE_PATH = Path(__file__).parents[1] / "scripts" / "present.py"
SPEC = importlib.util.spec_from_file_location("artifact_presenter", MODULE_PATH)
assert SPEC and SPEC.loader
MODULE = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = MODULE
SPEC.loader.exec_module(MODULE)
class FakeRunner:
def __init__(
self,
panes: list[dict[str, str]] | None = None,
*,
failed_run_panes: set[str] | None = None,
fail_rename: bool = False,
foreground_processes: dict[str, list[str]] | None = None,
shell_processes: dict[str, str] | None = None,
viewer_teardown_processes: dict[str, list[list[str]]] | None = None,
write_start_acknowledgment: bool = True,
list_delay_seconds: float = 0,
) -> None:
self.panes = panes or []
self.failed_run_panes = failed_run_panes or set()
self.fail_rename = fail_rename
self.foreground_processes = {
pane_id: list(processes)
for pane_id, processes in (foreground_processes or {}).items()
}
self.shell_processes = shell_processes or {}
self.viewer_teardown_processes = {
pane_id: [list(processes) for processes in snapshots]
for pane_id, snapshots in (viewer_teardown_processes or {}).items()
}
self.viewer_closing_panes: set[str] = set()
self.write_start_acknowledgment = write_start_acknowledgment
self.pending_commands: dict[str, str] = {}
self.list_delay_seconds = list_delay_seconds
self.calls: list[list[str]] = []
def __call__(self, command, **_kwargs):
argv = [str(value) for value in command]
self.calls.append(argv)
if argv[1:3] == ["pane", "send-text"]:
pane_id = argv[3]
if pane_id in self.failed_run_panes:
raise subprocess.CalledProcessError(1, argv)
self.pending_commands[pane_id] = argv[4]
command_parts = argv[4].strip().split()
decoded = [
base64.urlsafe_b64decode(value.encode("ascii")).decode("utf-8")
for value in command_parts[-4:]
]
if self.write_start_acknowledgment and decoded[-2] == "/usr/bin/mcat":
Path(decoded[-4]).write_text(decoded[-3], encoding="utf-8")
if argv[1:3] == ["pane", "send-keys"] and argv[4] == "q":
pane_id = argv[3]
self.viewer_closing_panes.add(pane_id)
if argv[1:3] == ["pane", "rename"] and self.fail_rename:
raise subprocess.CalledProcessError(1, argv)
if argv[1:3] == ["pane", "rename"]:
pane_id, label = argv[3], argv[4]
pane = next(
(candidate for candidate in self.panes if candidate["pane_id"] == pane_id),
None,
)
if pane is None:
self.panes.append({"pane_id": pane_id, "label": label})
else:
pane["label"] = label
if argv[1:3] == ["pane", "list"]:
pane_snapshot = [dict(pane) for pane in self.panes]
if self.list_delay_seconds:
time.sleep(self.list_delay_seconds)
output = {"result": {"panes": pane_snapshot}}
elif argv[1:3] == ["pane", "process-info"]:
pane_id = argv[4]
shell_name = self.shell_processes.get(pane_id, "fish")
teardown = self.viewer_teardown_processes.get(pane_id, [])
if pane_id in self.viewer_closing_panes and teardown:
process_names = teardown.pop(0)
elif pane_id in self.viewer_closing_panes:
process_names = [shell_name]
self.viewer_closing_panes.remove(pane_id)
self.foreground_processes[pane_id] = process_names
else:
process_names = self.foreground_processes.get(pane_id, [shell_name])
output = {
"result": {
"process_info": {
"pane_id": pane_id,
"shell_pid": 100,
"foreground_processes": [
{
"name": name,
"pid": 100 if name == shell_name else index + 101,
}
for index, name in enumerate(process_names)
],
}
}
}
elif argv[1:3] == ["pane", "split"]:
output = {"result": {"pane": {"pane_id": "workspace:p-new"}}}
else:
output = {}
return subprocess.CompletedProcess(argv, 0, json.dumps(output), "")
class OrcaAwareRunner(FakeRunner):
def __init__(self, *args, fail_orca_open: bool = False, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.fail_orca_open = fail_orca_open
def __call__(self, command, **kwargs):
argv = [str(value) for value in command]
if argv[1:3] == ["file", "open"]:
self.calls.append(argv)
if self.fail_orca_open:
raise subprocess.CalledProcessError(1, argv)
output = {"ok": True, "result": {"opened": True}}
return subprocess.CompletedProcess(argv, 0, json.dumps(output), "")
return super().__call__(command, **kwargs)
class ArtifactPresenterTests(unittest.TestCase):
def setUp(self) -> None:
self.tempdir = tempfile.TemporaryDirectory()
self.artifact = Path(self.tempdir.name) / "report.md"
self.artifact.write_text("# Report\n", encoding="utf-8")
self.env = {
"HERDR_WORKSPACE_ID": "workspace",
"HERDR_PANE_ID": "workspace:p-source",
}
self.executables = {
"herdr": "/usr/bin/herdr",
"mcat": "/usr/bin/mcat",
"orca": "/usr/bin/orca",
}
def tearDown(self) -> None:
self.tempdir.cleanup()
def which(self, name: str) -> str | None:
return self.executables.get(name)
def mcat_run_call(
self,
runner: FakeRunner,
pane_id: str,
artifact: Path | None = None,
) -> list[str]:
expected_artifact = str((artifact or self.artifact).resolve())
return next(
call
for call in runner.calls
if call[1:3] == ["pane", "send-text"]
and call[3] == pane_id
and [
base64.urlsafe_b64decode(value.encode("ascii")).decode("utf-8")
for value in call[4].strip().split()[-2:]
]
== ["/usr/bin/mcat", expected_artifact]
)
def test_reuses_named_pane(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
self.assertEqual(result.pane_id, "workspace:p-viewer")
self.mcat_run_call(runner, "workspace:p-viewer")
self.assertFalse(any(call[1:3] == ["pane", "split"] for call in runner.calls))
def test_closes_active_mcat_before_reusing_named_pane(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
foreground_processes={"workspace:p-viewer": ["mcat", "less"]},
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
close_viewer = [
"/usr/bin/herdr",
"pane",
"send-keys",
"workspace:p-viewer",
"q",
]
launch_viewer = self.mcat_run_call(runner, "workspace:p-viewer")
self.assertIn(close_viewer, runner.calls)
self.assertLess(runner.calls.index(close_viewer), runner.calls.index(launch_viewer))
def test_waits_through_transient_viewer_teardown_process(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
foreground_processes={"workspace:p-viewer": ["mcat", "less"]},
viewer_teardown_processes={"workspace:p-viewer": [["less"]]},
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
process_checks = [
call for call in runner.calls if call[1:3] == ["pane", "process-info"]
]
self.assertEqual(len(process_checks), 4)
self.mcat_run_call(runner, "workspace:p-viewer")
def test_requires_start_acknowledgment_after_launch(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
launch = self.mcat_run_call(runner, "workspace:p-viewer")
command = launch[4].strip().split()
decoded = [
base64.urlsafe_b64decode(value.encode("ascii")).decode("utf-8")
for value in command[-4:]
]
acknowledgment = Path(decoded[-4])
self.assertEqual(Path(command[1]).name, "launch-viewer.py")
self.assertEqual(acknowledgment.name, "viewer-started")
self.assertEqual(len(decoded[-3]), 24)
def test_does_not_report_success_without_start_acknowledgment(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
write_start_acknowledgment=False,
)
original_attempts = MODULE.VIEWER_START_ATTEMPTS
original_delay = MODULE.PANE_READY_DELAY_SECONDS
MODULE.VIEWER_START_ATTEMPTS = 2
MODULE.PANE_READY_DELAY_SECONDS = 0
try:
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
finally:
MODULE.VIEWER_START_ATTEMPTS = original_attempts
MODULE.PANE_READY_DELAY_SECONDS = original_delay
self.assertEqual(result.status, "artifact-only")
self.assertIn("could not initialize Herdr viewer", result.reason)
def test_accepts_idle_shell_by_pid_without_shell_name_allowlist(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
foreground_processes={"workspace:p-viewer": ["xonsh"]},
shell_processes={"workspace:p-viewer": "xonsh"},
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
self.assertFalse(
any(
call[1:3] == ["pane", "send-keys"] and call[4] == "q"
for call in runner.calls
)
)
def test_empty_foreground_state_is_not_treated_as_ready(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-transitioning", "label": MODULE.PANE_LABEL}],
foreground_processes={"workspace:p-transitioning": []},
)
original_attempts = MODULE.PANE_READY_ATTEMPTS
original_delay = MODULE.PANE_READY_DELAY_SECONDS
MODULE.PANE_READY_ATTEMPTS = 2
MODULE.PANE_READY_DELAY_SECONDS = 0
try:
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
finally:
MODULE.PANE_READY_ATTEMPTS = original_attempts
MODULE.PANE_READY_DELAY_SECONDS = original_delay
self.assertEqual(result.status, "opened")
process_checks = [
call
for call in runner.calls
if call[1:3] == ["pane", "process-info"]
and call[4] == "workspace:p-transitioning"
]
self.assertEqual(len(process_checks), 2)
def test_does_not_interrupt_unrelated_foreground_process(self) -> None:
runner = FakeRunner(
[{"pane_id": "workspace:p-busy", "label": MODULE.PANE_LABEL}],
foreground_processes={"workspace:p-busy": ["vim"]},
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "opened")
self.assertFalse(
any(
call[1:3] == ["pane", "send-keys"] and call[4] == "q"
for call in runner.calls
)
)
self.assertIn(
[
"/usr/bin/herdr",
"pane",
"rename",
"workspace:p-busy",
f"{MODULE.PANE_LABEL} stale",
],
runner.calls,
)
def test_tries_each_named_pane_before_creating_another(self) -> None:
runner = FakeRunner(
[
{"pane_id": "workspace:p-stale", "label": MODULE.PANE_LABEL},
{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL},
],
failed_run_panes={"workspace:p-stale"},
)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
self.assertEqual(result.pane_id, "workspace:p-viewer")
self.assertIn(
[
"/usr/bin/herdr",
"pane",
"rename",
"workspace:p-stale",
f"{MODULE.PANE_LABEL} stale",
],
runner.calls,
)
self.assertFalse(any(call[1:3] == ["pane", "split"] for call in runner.calls))
def test_opens_editor_friendly_file_in_orca_managed_session(self) -> None:
source = self.artifact.with_name("module.ts")
source.write_text("export const answer = 42;\n", encoding="utf-8")
runner = OrcaAwareRunner()
result = MODULE.present_document(
source,
env={**self.env, "ORCA_WORKSPACE_ID": "workspace"},
which=self.which,
run=runner,
)
self.assertEqual(result.status, "orca-opened")
self.assertEqual(
runner.calls,
[["/usr/bin/orca", "file", "open", str(source.resolve()), "--json"]],
)
def test_uses_configured_orca_cli_command(self) -> None:
self.executables["orca-wsl"] = "/opt/orca/orca-wsl"
runner = OrcaAwareRunner()
result = MODULE.present_document(
self.artifact,
env={"TERM_PROGRAM": "Orca", "ORCA_CLI_COMMAND": "orca-wsl"},
which=self.which,
run=runner,
)
self.assertEqual(result.status, "orca-opened")
self.assertEqual(runner.calls[0][0], "/opt/orca/orca-wsl")
def test_falls_back_to_mcat_when_orca_cannot_open_text_file(self) -> None:
runner = OrcaAwareRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
fail_orca_open=True,
)
result = MODULE.present_document(
self.artifact,
env={**self.env, "ORCA_TERMINAL_HANDLE": "term-1"},
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
self.assertEqual(runner.calls[0][1:3], ["file", "open"])
self.mcat_run_call(runner, "workspace:p-viewer")
def test_presents_binary_file_directly_with_mcat_in_orca_session(self) -> None:
video = self.artifact.with_name("review.mp4")
video.write_bytes(b"video\0fixture")
runner = OrcaAwareRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
)
result = MODULE.present_document(
video,
env={**self.env, "TERM_PROGRAM": "Orca"},
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
self.assertFalse(any(call[1:3] == ["file", "open"] for call in runner.calls))
self.mcat_run_call(runner, "workspace:p-viewer", video)
def test_encodes_artifact_paths_before_terminal_delivery(self) -> None:
artifact = self.artifact.with_name("report '$(touch unsafe)'.md")
artifact.write_text("# Safe\n", encoding="utf-8")
runner = FakeRunner(
[{"pane_id": "workspace:p-viewer", "label": MODULE.PANE_LABEL}],
)
result = MODULE.present_document(
artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "reused")
launch = self.mcat_run_call(runner, "workspace:p-viewer", artifact)
command_parts = launch[4].strip().split()
self.assertTrue(all(MODULE.SAFE_COMMAND_TOKEN.fullmatch(part) for part in command_parts))
self.assertNotIn("$(touch unsafe)", launch[4])
def test_creates_right_pane_when_viewer_is_absent(self) -> None:
runner = FakeRunner()
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "opened")
split = next(call for call in runner.calls if call[1:3] == ["pane", "split"])
self.assertEqual(
split,
[
"/usr/bin/herdr",
"pane",
"split",
"--pane",
"workspace:p-source",
"--direction",
"right",
"--cwd",
str(self.artifact.parent.resolve()),
"--no-focus",
],
)
self.assertIn(
[
"/usr/bin/herdr",
"pane",
"rename",
"workspace:p-new",
MODULE.PANE_LABEL,
],
runner.calls,
)
def test_closes_new_pane_when_rename_fails(self) -> None:
runner = FakeRunner(fail_rename=True)
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
self.assertEqual(result.status, "artifact-only")
self.assertIn("could not initialize Herdr viewer", result.reason)
self.assertIn(
["/usr/bin/herdr", "pane", "close", "workspace:p-new"],
runner.calls,
)
def test_falls_back_to_artifact_when_mcat_is_unavailable(self) -> None:
self.executables.pop("mcat")
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=FakeRunner(),
)
self.assertEqual(result.status, "artifact-only")
self.assertEqual(result.reason, "mcat is unavailable")
def test_pane_lock_serializes_concurrent_presenters(self) -> None:
first_acquired = threading.Event()
release_first = threading.Event()
second_acquired = threading.Event()
original_runtime = os.environ.get("XDG_RUNTIME_DIR")
os.environ["XDG_RUNTIME_DIR"] = self.tempdir.name
def first_presenter() -> None:
with MODULE.lock_pane("workspace:p-viewer"):
first_acquired.set()
release_first.wait(timeout=1)
def second_presenter() -> None:
first_acquired.wait(timeout=1)
with MODULE.lock_pane("workspace:p-viewer"):
second_acquired.set()
first = threading.Thread(target=first_presenter)
second = threading.Thread(target=second_presenter)
try:
first.start()
second.start()
self.assertTrue(first_acquired.wait(timeout=1))
self.assertFalse(second_acquired.wait(timeout=0.05))
release_first.set()
self.assertTrue(second_acquired.wait(timeout=1))
finally:
release_first.set()
first.join(timeout=1)
second.join(timeout=1)
if original_runtime is None:
os.environ.pop("XDG_RUNTIME_DIR", None)
else:
os.environ["XDG_RUNTIME_DIR"] = original_runtime
def test_workspace_lock_prevents_duplicate_viewer_creation(self) -> None:
runner = FakeRunner(list_delay_seconds=0.05)
results: list[MODULE.PresentationResult] = []
results_lock = threading.Lock()
original_runtime = os.environ.get("XDG_RUNTIME_DIR")
os.environ["XDG_RUNTIME_DIR"] = self.tempdir.name
def present() -> None:
result = MODULE.present_document(
self.artifact,
env=self.env,
which=self.which,
run=runner,
)
with results_lock:
results.append(result)
threads = [threading.Thread(target=present) for _ in range(2)]
try:
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=2)
finally:
if original_runtime is None:
os.environ.pop("XDG_RUNTIME_DIR", None)
else:
os.environ["XDG_RUNTIME_DIR"] = original_runtime
self.assertEqual(sorted(result.status for result in results), ["opened", "reused"])
splits = [call for call in runner.calls if call[1:3] == ["pane", "split"]]
self.assertEqual(len(splits), 1)
def test_rejects_missing_artifact(self) -> None:
with self.assertRaisesRegex(ValueError, "artifact is not a file"):
MODULE.present_document(
self.artifact.with_name("missing.md"),
env=self.env,
which=self.which,
run=FakeRunner(),
)
if __name__ == "__main__":
unittest.main()