rs-mrxvt/scripts/stress_test.py

#!/usr/bin/env python3
"""
rs-mrxvt stress harness — spawns N rs-mrxvt tabs in broadcast mode and verifies
that input typed once is mirrored to every tab's PTY.

This is the "test_suite.py" referenced at the end of the original design chat.
It uses only stdlib + the `mrxvt` library's public API via cargo-test FFI is
NOT used — instead, we spawn the compiled `rs-mrxvt` binary with `--broadcast`
and verify behavior through PTY inspection.

Since the GUI binary can't be driven headlessly here, this script tests the
underlying PTY plumbing directly by spawning subprocesses through Python's
pty module and verifying broadcasting at the shell level. It's a smoke test
for the "could I manage 50 servers simultaneously?" use case from the chat.

Usage:
    python3 scripts/stress_test.py [--tabs 50] [--timeout 30]

Requirements:
    - Python 3.10+
    - /bin/sh
    - (Optionally) rs-mrxvt binary for an integration smoke check
"""

from __future__ import annotations

import argparse
import os
import pty
import select
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import List


DEFAULT_TABS = 50
DEFAULT_TIMEOUT = 30  # seconds
MARKER = "BROADCAST_STRESS_MARKER_42"


@dataclass
class TabSession:
    """One shell-in-a-PTY."""
    master_fd: int
    pid: int
    title: str
    buffer: bytes = b""

    def write(self, data: bytes) -> None:
        os.write(self.master_fd, data)

    def read_nonblock(self, max_bytes: int = 4096) -> bytes:
        try:
            ready, _, _ = select.select([self.master_fd], [], [], 0.05)
            if ready:
                chunk = os.read(self.master_fd, max_bytes)
                self.buffer += chunk
                return chunk
        except OSError:
            pass
        return b""

    def has_marker(self, marker: str) -> bool:
        return marker.encode() in self.buffer

    def close(self) -> None:
        try:
            os.close(self.master_fd)
        except OSError:
            pass
        try:
            os.kill(self.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass


def spawn_tab(title: str) -> TabSession:
    """Spawn /bin/sh in a new PTY, returning the master FD and PID."""
    pid, master_fd = pty.fork()
    if pid == 0:
        # Child
        os.environ["TERM"] = "xterm-256color"
        os.execvp("/bin/sh", ["/bin/sh", "-i"])
    return TabSession(master_fd=master_fd, pid=pid, title=title)


def drive_broadcast(tabs: List[TabSession], marker: str, timeout: float) -> bool:
    """Send `marker` to ONE tab (simulating broadcast mode by writing to all),
    then verify every tab received it via shell echo.

    In rs-mrxvt's broadcast-all mode, a single keystroke goes to all PTYs.
    Here we simulate by sending the marker to every tab's stdin directly.
    The test is: can 50 PTYs all receive the same input within `timeout`?
    """
    deadline = time.time() + timeout

    # Send marker + newline to every tab.
    payload = marker.encode() + b"\n"
    for tab in tabs:
        tab.write(payload)

    # Poll every tab until each has the marker in its output.
    pending = list(tabs)
    while pending and time.time() < deadline:
        new_pending = []
        for tab in pending:
            tab.read_nonblock()
            if tab.has_marker(marker):
                continue  # done
            new_pending.append(tab)
        pending = new_pending
        time.sleep(0.02)

    return not pending  # True if all tabs saw the marker


def smoke_test_binary(binary: Path) -> bool:
    """If rs-mrxvt binary exists, ensure it at least --helps without crashing."""
    if not binary.exists():
        print(f"  (skipping binary smoke test: {binary} not found)")
        return True
    try:
        proc = subprocess.run(
            [str(binary), "--help"],
            capture_output=True,
            timeout=5,
        )
        if proc.returncode == 0:
            print(f"  binary --help OK ({binary})")
            return True
        print(f"  binary --help returned {proc.returncode}", file=sys.stderr)
        return False
    except subprocess.TimeoutExpired:
        print(f"  binary --help timed out", file=sys.stderr)
        return False
    except FileNotFoundError:
        return True


def main() -> int:
    parser = argparse.ArgumentParser(
        description="rs-mrxvt broadcast stress harness",
    )
    parser.add_argument("--tabs", type=int, default=DEFAULT_TABS,
                        help=f"number of concurrent tabs (default: {DEFAULT_TABS})")
    parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT,
                        help=f"per-phase timeout in seconds (default: {DEFAULT_TIMEOUT})")
    parser.add_argument("--binary", type=Path,
                        default=Path("target/release/rs-mrxvt"),
                        help="path to rs-mrxvt binary for smoke test")
    args = parser.parse_args()

    print(f"=== rs-mrxvt stress harness ===")
    print(f"tabs: {args.tabs}")
    print(f"timeout: {args.timeout}s")
    print()

    # Phase 0: binary smoke test
    print("[phase 0] binary smoke test")
    if not smoke_test_binary(args.binary):
        return 1

    # Phase 1: spawn N tabs
    print(f"[phase 1] spawning {args.tabs} PTY-backed shells...")
    tabs: List[TabSession] = []
    t0 = time.time()
    for i in range(args.tabs):
        try:
            tab = spawn_tab(f"tab-{i+1}")
            tabs.append(tab)
        except OSError as e:
            print(f"  failed to spawn tab {i+1}: {e}", file=sys.stderr)
            break
    spawn_time = time.time() - t0
    print(f"  spawned {len(tabs)} tabs in {spawn_time:.2f}s")

    if len(tabs) != args.tabs:
        print(f"  FAIL: only spawned {len(tabs)}/{args.tabs} tabs", file=sys.stderr)
        for tab in tabs:
            tab.close()
        return 1

    # Phase 2: broadcast marker
    print(f"[phase 2] broadcasting marker to all {len(tabs)} tabs...")
    t0 = time.time()
    ok = drive_broadcast(tabs, MARKER, args.timeout)
    elapsed = time.time() - t0
    if not ok:
        # Find which tabs missed it
        missing = [i for i, t in enumerate(tabs) if not t.has_marker(MARKER)]
        print(f"  FAIL: {len(missing)} tabs did not receive marker: {missing[:10]}{'...' if len(missing) > 10 else ''}",
              file=sys.stderr)
        for tab in tabs:
            tab.close()
        return 1
    print(f"  OK: all {len(tabs)} tabs received marker in {elapsed:.2f}s "
          f"({len(tabs)/elapsed:.1f} tabs/sec)")

    # Phase 3: cleanup
    print("[phase 3] cleanup")
    for tab in tabs:
        tab.close()
    print(f"  closed {len(tabs)} tabs")

    print()
    print(f"=== PASS: {args.tabs}-tab broadcast stress test passed ===")
    return 0


if __name__ == "__main__":
    sys.exit(main())