tai/taicpp/tests/foundation_fixtures.py

"""Process fixtures for tests/foundation_parity.py: a scripted native-plexer
backend on a unix socket and an isolated tmux server."""
import json
import os
import shutil
import socket
import subprocess
import threading


class FakeBackend:
    """Answers the NativePlexerBackend line protocol and logs every request.

    send/read/status answer fixed data; show/result run a two-poll flow whose
    outcome follows the screen title (cancel, timeout, weird, error, else ok);
    cmd "garbage" answers non-JSON and cmd "refuse" answers ok=false."""

    def __init__(self, path):
        self.path = path
        self.log = []
        self.flows = {}
        self.counter = 0
        self.server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.server.bind(path)
        self.server.listen(8)
        self.thread = threading.Thread(target=self.serve, daemon=True)
        self.thread.start()

    def serve(self):
        while True:
            try:
                conn, _ = self.server.accept()
            except OSError:
                return
            with conn:
                data = b""
                while not data.endswith(b"\n"):
                    chunk = conn.recv(65536)
                    if not chunk:
                        break
                    data += chunk
                try:
                    request = json.loads(data.decode("utf-8"))
                except ValueError:
                    continue
                self.log.append(request)
                conn.sendall(self.answer(request))

    def answer(self, request):
        cmd = request.get("cmd")
        if cmd == "send":
            return self.ok(True)
        if cmd == "read":
            return self.ok("screen of %s lines\nsecond line ü" % request.get("lines"))
        if cmd == "status":
            return self.ok({"backend": "native", "pane": request.get("pane")})
        if cmd == "garbage":
            return b"this is not json\n"
        if cmd == "refuse":
            return (json.dumps({"ok": False}) + "\n").encode()
        if cmd == "boom":
            return (json.dumps({"ok": False, "error": "boom"}) + "\n").encode()
        if cmd == "show":
            self.counter += 1
            flow_id = "f%d" % self.counter
            self.flows[flow_id] = {"screen": request.get("screen") or {}, "polls": 0}
            return self.ok({"id": flow_id})
        if cmd == "result":
            flow = self.flows.get(request.get("id"))
            if flow is None:
                return (json.dumps({"ok": False, "error": "unknown flow"}) + "\n").encode()
            flow["polls"] += 1
            if flow["polls"] < 2:
                return self.ok({"done": False})
            title = flow["screen"].get("title")
            if title == "cancel":
                return self.ok({"done": True, "status": "cancelled"})
            if title == "timeout":
                return self.ok({"done": True, "status": "timeout"})
            if title == "weird":
                return self.ok({"done": True, "status": "weird"})
            if title == "never":
                return self.ok({"done": False})
            return self.ok({"done": True, "status": "ok", "value": 1 if flow["screen"].get("kind") == "choice" else "typed"})
        if cmd == "cancel":
            return self.ok(True)
        return (json.dumps({"ok": False, "error": "unknown cmd %s" % cmd}) + "\n").encode()

    @staticmethod
    def ok(data):
        return (json.dumps({"ok": True, "data": data}) + "\n").encode()

    def close(self):
        try:
            self.server.close()
        finally:
            if os.path.exists(self.path):
                os.unlink(self.path)


class AttachedClient:
    """A tmux client attached through a pty, so display-menu and display-popup have a client to draw on.
    The client's drawing is collected so a test can wait for a title to appear before it types."""

    def __init__(self, server, session="t", cols=120, rows=40):
        import fcntl
        import pty
        import struct
        import termios
        self.output = b""
        self.lock = threading.Lock()
        pid, fd = pty.fork()
        if pid == 0:
            env = dict(server.env, TERM="xterm-256color")
            os.execvpe("tmux", ["tmux", "-f", "/dev/null", "attach-session", "-t", session], env)
        self.pid, self.fd = pid, fd
        fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
        self.thread = threading.Thread(target=self.pump, daemon=True)
        self.thread.start()

    def pump(self):
        while True:
            try:
                chunk = os.read(self.fd, 65536)
            except OSError:
                return
            if not chunk:
                return
            with self.lock:
                self.output += chunk

    def mark(self):
        with self.lock:
            return len(self.output)

    def wait_for(self, needle, since=0, seconds=15.0):
        import time
        end = time.time() + seconds
        target = needle.encode()
        while time.time() < end:
            with self.lock:
                if target in self.output[since:]:
                    return True
            time.sleep(0.05)
        return False

    def type(self, data):
        os.write(self.fd, data if isinstance(data, bytes) else data.encode())

    def close(self):
        try:
            os.close(self.fd)
        except OSError:
            pass
        try:
            os.kill(self.pid, 9)
            os.waitpid(self.pid, 0)
        except OSError:
            pass


class TmuxServer:
    """A private tmux server (own TMUX_TMPDIR) with one session running `cat`, so typed text is echoed back."""

    def __init__(self, root):
        self.tmpdir = os.path.join(root, "tmux")
        os.makedirs(self.tmpdir)
        self.env = {key: value for key, value in os.environ.items() if not key.startswith("TMUX")}
        self.env["TMUX_TMPDIR"] = self.tmpdir
        self.available = shutil.which("tmux") is not None
        if self.available:
            subprocess.run(["tmux", "-f", "/dev/null", "new-session", "-d", "-s", "t", "-x", "80", "-y", "24", "cat"], env=self.env, check=True,
                           stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    def close(self):
        if self.available:
            subprocess.run(["tmux", "kill-server"], env=self.env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)