tai/taicpp/tests/repl_keys.py

#!/usr/bin/env python3
"""Ctrl-key behavior of the detz REPL over a real pty (no model needed).

Covers the shared-stdin fix: Ctrl-C once re-prompts, Ctrl-C during a secret
prompt aborts it and the next prompt still reads (no 0 ms phantom answer),
double Ctrl-C exits 0, Ctrl-D exits 0, Ctrl-Z suspends and SIGCONT resumes.

usage: python3 -I tests/repl_keys.py build/detz
"""
import os
import pty
import re
import select
import signal
import subprocess
import sys
import tempfile
import time

DET = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else "build/detz")
CHECKS = 0


def check(cond, label):
    global CHECKS
    CHECKS += 1
    if not cond:
        raise AssertionError("FAILED: " + label)
    print("ok %d: %s" % (CHECKS, label))


class Proc:
    def __init__(self, pid):
        self.pid = pid
        self.returncode = None

    def poll(self):
        if self.returncode is not None:
            return self.returncode
        pid, status = os.waitpid(self.pid, os.WNOHANG)
        if pid == 0:
            return None
        self.returncode = os.waitstatus_to_exitcode(status)
        return self.returncode

    def wait(self, timeout=15):
        end = time.time() + timeout
        while time.time() < end:
            code = self.poll()
            if code is not None:
                return code
            time.sleep(0.05)
        raise subprocess.TimeoutExpired(self.pid, timeout)

    def kill(self):
        try:
            os.kill(self.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


def spawn(home):
    pid, fd = pty.fork()
    if pid == 0:
        os.environ["TAI_HOME"] = home
        os.environ["TERM"] = "xterm"
        os.environ.pop("TMUX", None)
        os.environ.pop("TMUX_PANE", None)
        os.chdir(home)
        os.execv(DET, [DET])
        os._exit(127)
    return Proc(pid), fd


def spawn_keeplink(home):
    """Spawn detz in its own foreground process group with a middle process
    keeping the session: without it the group is orphaned and the kernel
    discards SIGTSTP, so no Ctrl-Z test could pass."""
    pid, fd = pty.fork()
    if pid == 0:
        grandchild = os.fork()
        if grandchild == 0:
            os.setpgid(0, 0)
            os.environ["TAI_HOME"] = home
            os.environ["TERM"] = "xterm"
            os.environ.pop("TMUX", None)
            os.environ.pop("TMUX_PANE", None)
            os.chdir(home)
            os.execv(DET, [DET])
            os._exit(127)
        os.tcsetpgrp(0, grandchild)
        print("grandchild %d" % grandchild, flush=True)
        while True:
            try:
                done = os.waitpid(grandchild, 0)
            except ChildProcessError:
                break
            print("reaped %r" % (done,), flush=True)
            break
        os._exit(0)
    return Proc(pid), fd


def proc_state(pid):
    with open("/proc/%d/status" % pid) as handle:
        for line in handle:
            if line.startswith("State:"):
                return line.split()[1]
    return "?"


def read_until(fd, markers, timeout=30):
    if isinstance(markers, str):
        markers = [markers]
    buf = b""
    end = time.time() + timeout
    while time.time() < end:
        ready, _, _ = select.select([fd], [], [], max(0, end - time.time()))
        if not ready:
            break
        try:
            chunk = os.read(fd, 65536)
        except OSError:
            break
        if not chunk:
            break
        buf += chunk
        text = buf.decode("utf-8", "replace")
        for marker in markers:
            if marker in text:
                return text
    return buf.decode("utf-8", "replace")


def drain(fd, seconds=0.5):
    end = time.time() + seconds
    while time.time() < end:
        ready, _, _ = select.select([fd], [], [], max(0, end - time.time()))
        if not ready:
            break
        try:
            if not os.read(fd, 65536):
                break
        except OSError:
            break


def wait_exit(proc, timeout=15):
    try:
        return proc.wait(timeout=timeout)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait()
        raise AssertionError("process did not exit")


def main():
    home = tempfile.mkdtemp(prefix="repl-keys-")
    # 1: Ctrl-C once re-prompts, process stays alive
    proc, fd = spawn(home)
    out = read_until(fd, "type /help for commands")
    check("type /help for commands" in out, "banner shown")
    out = read_until(fd, "›")
    check("›" in out, "prompt shown")
    os.write(fd, b"\x03")
    out = read_until(fd, "›")
    check("›" in out and proc.poll() is None, "single Ctrl-C re-prompts, stays alive")
    # 2: Ctrl-C aborts a secret prompt, next prompt still reads
    os.write(fd, b"/secret set probe1\n")
    out = read_until(fd, "value for 'probe1'")
    check("value for 'probe1'" in out, "secret prompt shown")
    os.write(fd, b"\x03")
    out = read_until(fd, ["aborted", "›"])
    check("aborted" in out, "Ctrl-C aborts secret prompt")
    read_until(fd, "›")
    os.write(fd, b"/secret set probe1\n")
    out = read_until(fd, "value for 'probe1'")
    check("value for 'probe1'" in out, "secret prompt shown again")
    time.sleep(0.3)
    os.write(fd, b"s3cr3t-value\n")
    out = read_until(fd, ["stored secret 'probe1'", "›"], timeout=15)
    check("stored secret 'probe1'" in out, "answer after Ctrl-C is read (no phantom empty)")
    read_until(fd, "›")
    # 3: double Ctrl-C exits 0 with a saved session (two distinct keypresses:
    # the tty discards a second ^C written in the same instant as the first)
    os.write(fd, b"\x03")
    read_until(fd, "›")
    os.write(fd, b"\x03")
    code = wait_exit(proc)
    check(code == 0, "double Ctrl-C exits 0 (got %r)" % code)
    os.close(fd)
    # 4: Ctrl-D exits 0
    proc, fd = spawn(home)
    read_until(fd, "›")
    os.write(fd, b"\x04")
    code = wait_exit(proc)
    check(code == 0, "Ctrl-D exits 0 (got %r)" % code)
    os.close(fd)
    # 5: Ctrl-Z suspends, SIGCONT resumes, /quit exits
    proc, fd = spawn_keeplink(home)
    out = read_until(fd, "grandchild ")
    grandchild = int(re.search(r"grandchild (\d+)", out).group(1))
    read_until(fd, "›")
    os.write(fd, b"\x1a")
    time.sleep(1.0)
    check(proc_state(grandchild) == "T", "Ctrl-Z suspends the process")
    os.kill(grandchild, signal.SIGCONT)
    time.sleep(1.0)
    check(proc_state(grandchild) in ("S", "R"), "SIGCONT resumes the process")
    drain(fd)
    os.write(fd, b"/quit\n")
    code = wait_exit(proc)
    check(code == 0, "/quit after resume exits 0 (got %r)" % code)
    os.close(fd)
    # vault state, not rendering
    prof = os.path.join(home, "profiles")
    sessions = [f for f in os.listdir(prof) if ".session." in f and f.endswith(".json")]
    check(len(sessions) >= 1, "session file saved (%d found)" % len(sessions))
    print("%d checks passed" % CHECKS)


main()