cont/tests/pty_smoke.py
2026-09-15 14:15:06 +01:00

230 lines
9.6 KiB
Python

#!/usr/bin/env python3
"""Unix PTY integration checks. No real agents or user session files are touched."""
import fcntl
import json
import os
from pathlib import Path
import pty
import re
import unicodedata
import select
import signal
import struct
import sqlite3
import subprocess
import sys
import tempfile
import termios
import time
BINARY = str(Path(sys.argv[1] if len(sys.argv) > 1 else "target/release/cont").resolve())
class Picker:
def __init__(self, args, env, cwd):
self.master, self.slave = pty.openpty()
self.resize(120, 30)
self.before = termios.tcgetattr(self.slave)
self.output = b""
self.history = b""
def setup():
os.setsid()
fcntl.ioctl(0, termios.TIOCSCTTY, 0)
self.proc = subprocess.Popen([BINARY, *args], stdin=self.slave, stdout=self.slave,
stderr=self.slave, env=env, cwd=cwd, preexec_fn=setup)
self.wait_for(b"Ctrl-R refresh")
def resize(self, width, height):
self.output = b""
self.width, self.height = width, height
fcntl.ioctl(self.slave, termios.TIOCSWINSZ, struct.pack("HHHH", height, width, 0, 0))
if hasattr(self, "proc"):
os.kill(self.proc.pid, signal.SIGWINCH)
def read(self, timeout=0.05):
if select.select([self.master], [], [], timeout)[0]:
try:
data = os.read(self.master, 65536)
self.output += data
self.history += data
except OSError:
pass
def screen(self):
# Minimal emulator for Ratatui's absolute-cursor/SGR output. Reconstruct
# screen cells because later frames emit diffs, not complete lines.
cells = [[" "] * self.width for _ in range(self.height)]
x = y = 0
for match in re.finditer(r"\x1b\[([0-?]*)([ -/]*)([@-~])|([^\x1b])", self.history.decode(errors="replace"), re.S):
params, _, command, char = match.groups()
if command in ("H", "f"):
parts = params.split(";")
y = int(parts[0] or 1) - 1
x = int(parts[1] or 1) - 1 if len(parts) > 1 else 0
elif command == "J" and params == "2":
cells = [[" "] * self.width for _ in range(self.height)]
elif char == "\r":
x = 0
elif char == "\n":
y += 1
elif char and ord(char) >= 32:
wide = unicodedata.east_asian_width(char) in ("W", "F")
if 0 <= x < self.width and 0 <= y < self.height:
cells[y][x] = char
if wide and x + 1 < self.width:
cells[y][x + 1] = ""
x += 2 if wide else 1
return "\n".join("".join(row) for row in cells)
def wait_for(self, marker):
end = time.monotonic() + 8
self.read()
while marker not in self.output and marker.decode() not in self.screen() and time.monotonic() < end:
self.read()
assert marker in self.output or marker.decode() in self.screen(), (marker, self.screen())
def send(self, keys):
self.output = b""
os.write(self.master, keys)
def finish(self, code):
end = time.monotonic() + 8
while self.proc.poll() is None and time.monotonic() < end:
self.read()
if self.proc.poll() is None:
self.proc.kill()
raise AssertionError(("Picker did not exit", self.screen(), self.output[-2000:]))
self.read()
assert self.proc.returncode == code, (self.proc.returncode, self.output[-4000:])
after = termios.tcgetattr(self.master)
assert after == self.before, "Terminal settings were not restored"
assert b"\x1b[?1049l" in self.output, "Alternate screen was not left"
os.close(self.master)
os.close(self.slave)
def main():
with tempfile.TemporaryDirectory(prefix="cont-test-") as tmp:
root = Path(tmp)
work = root / "work with 'quotes'"
other = root / "other"
for p in [work, other, root / "pi", root / "claude/projects/p", root / "codex/sessions", root / ".omp/agent/sessions", root / "data/opencode", root / "bin"]:
p.mkdir(parents=True)
env = os.environ | {
"HOME": str(root),
"XDG_DATA_HOME": str(root / "data"),
"OPENCODE_DB": str(root / "data/opencode/opencode.db"),
"PI_CODING_AGENT_SESSION_DIR": str(root / "pi"),
"CLAUDE_CONFIG_DIR": str(root / "claude"),
"CODEX_HOME": str(root / "codex"),
"RESUME_CACHE_DIR": str(root / "cache"),
"PATH": str(root / "bin") + os.pathsep + os.environ["PATH"],
"RESUME_TEST_LOG": str(root / "handoff.json"),
"TERM": "xterm-256color",
}
sessions = {
".omp/agent/sessions/omp.jsonl": [
{"type": "title", "title": "OMP needle"},
{"type": "session", "id": "omp-id", "cwd": str(other)},
{"type": "message", "message": {"role": "user", "content": "omp prompt"}},
],
"pi/a.jsonl": [
{"type": "session", "id": "pi-id", "cwd": str(work)},
{"type": "session_info", "name": "Pi needle"},
{"type": "message", "message": {"role": "user", "content": "日本語 café"}},
],
"claude/projects/p/c.jsonl": [
{"type": "user", "sessionId": "claude-id", "cwd": str(other), "message": {"role": "user", "content": "Claude needle"}},
],
"codex/sessions/x.jsonl": [
{"type": "session_meta", "payload": {"id": "codex-id", "cwd": str(work)}},
{"type": "event_msg", "payload": {"type": "user_message", "message": "Codex needle"}},
],
}
for path, records in sessions.items():
(root / path).write_text("".join(json.dumps(r) + "\n" for r in records))
database = root / "data/opencode/opencode.db"
conn = sqlite3.connect(database)
conn.executescript("""
CREATE TABLE session (id TEXT, directory TEXT, title TEXT, time_updated INTEGER, parent_id TEXT, time_archived INTEGER);
CREATE TABLE message (id TEXT, session_id TEXT, time_created INTEGER, data TEXT);
CREATE TABLE part (id TEXT, message_id TEXT, session_id TEXT, time_created INTEGER, data TEXT);
""")
conn.execute("INSERT INTO session VALUES ('ses_oc', ?, 'OpenCode needle', 1000, NULL, NULL)", [str(other)])
conn.commit()
conn.close()
stub = f"#!{sys.executable}\n" + '''import json, os, sys, termios
flags = termios.tcgetattr(0)[3]
with open(os.environ["RESUME_TEST_LOG"], "w") as f:
json.dump({"agent": os.path.basename(sys.argv[0]), "args": sys.argv[1:], "cwd": os.getcwd(),
"canonical": bool(flags & termios.ICANON), "echo": bool(flags & termios.ECHO),
"database": os.environ.get("OPENCODE_DB")}, f)
sys.exit(23)
'''
for agent in ["pi", "claude", "codex", "opencode", "omp"]:
path = root / "bin" / agent
path.write_text(stub)
path.chmod(0o755)
# Actual exec handoff, argv boundaries, original cwd, exit status, terminal flags.
for agent, args, cwd in [
("pi", ["--session", str((root / "pi/a.jsonl").resolve())], work),
("claude", ["--resume", "claude-id"], other),
("codex", ["resume", "codex-id"], work),
("opencode", ["--session", "ses_oc"], other),
("omp", ["--resume", str((root / ".omp/agent/sessions/omp.jsonl").resolve())], other),
]:
p = Picker(["--all", "--agent", agent], env, work)
p.send(b"\r")
p.finish(23)
log = json.loads((root / "handoff.json").read_text())
selected_database = log.pop("database")
if agent == "opencode":
assert selected_database == str(database.resolve())
assert log == {"agent": agent, "args": args, "cwd": str(cwd.resolve()), "canonical": True, "echo": True}, log
# Scope, agent cycling, Unicode paste, clearing, location selection, resize, quit.
p = Picker([], env, work)
assert "2 results" in p.screen(), p.screen()
p.send(b"\t")
p.wait_for(b"All locations")
p.send(b"\x1b[Z")
p.wait_for(b"1 results")
p.send(b"\x1b[200~" + "日本語".encode() + b"\x1b[201~")
p.wait_for("日本語".encode())
p.send(b"\x15") # Ctrl-U
p.wait_for(b"All locations")
p.send(b"\x1b[Z") # Claude
p.wait_for(b"Claude needle")
p.send(b"\x0c") # Ctrl-L
p.wait_for(b"Locations")
p.send(b"other")
p.wait_for(b"other")
p.send(b"\r")
p.wait_for(b"1 results")
p.resize(30, 8)
p.wait_for(b"Enlarge terminal")
p.resize(120, 30)
p.wait_for(b"Enter resume")
p.send(b"\x03")
p.finish(0)
# Print-only selection and failed executable must also restore the terminal.
p = Picker(["--all", "--agent", "pi", "--print"], env, work)
p.send(b"\r")
p.finish(0)
assert b"cd -- '" in p.output
assert b"'--session'" in p.output
(root / "bin/pi").unlink()
env["PATH"] = str(root / "bin") # Deliberately no fallback to the real Pi.
p = Picker(["--all", "--agent", "pi"], env, work)
p.send(b"\r")
p.finish(1)
assert b"Could not start pi" in p.output
print("PTY smoke tests passed: all agents, scope/filter/paste/resize, print, errors, terminal restoration")
if __name__ == "__main__":
main()