tai/tests/grr_tui.py

#!/usr/bin/env python3
"""tests/grr_tui.py - integration tests for `tai.py --tui`, the terminal client of the grr server.

Built on tai_test.py (see tests/README.md). Each case starts a real grr server
from the sibling `grr` checkout on a throwaway GRR_HOME and port, signs in with
the DevPlace key of this machine (DEVPLACE_API_KEY, the same key the real server
uses to reach the gateway), and drives the real `tai.py --tui` over a pty with
real key bytes and real mouse reports. Nothing is mocked: the bot that answers is
a real model behind the real gateway, so the LLM-driven case is probabilistic
(a free, rotating backend) and says so in its failure text.

What is checked is state, not rendering: the server's own API answers whether the
message was stored, whether the reply exists, whether the setting changed. The
screen is only read for the one thing a screen is for, that the user can see it
(the typed words and the reply appear in the output).

Skipped loudly (never silently) when the grr checkout, DEVPLACE_API_KEY or a free
port is missing.
"""
import asyncio
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from tai_test import TestCase, TestAssertionFailed, TestTimeout, main_cli, strip_ansi  # noqa: E402

GRR_DIR = Path(os.environ.get("GRR_DIR") or Path(__file__).resolve().parent.parent.parent / "grr")


def free_port():
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]


class GrrServer:
    """A throwaway grr server: own home, own ports, killed with its process group."""

    def __init__(self):
        self.home = tempfile.mkdtemp(prefix="grr-it-")
        self.port = free_port()
        self.url = "http://127.0.0.1:%d" % self.port
        self.proc = None
        self.token = ""

    def start(self):
        env = dict(os.environ, GRR_HOME=self.home, GRR_APPS_PORT=str(free_port()), PYTHONPATH=str(GRR_DIR))
        self.proc = subprocess.Popen([sys.executable, "-m", "grr", "serve", "--host", "127.0.0.1", "--port", str(self.port), "--no-reload"],
                                     cwd=str(GRR_DIR), env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
        deadline = time.time() + 30
        while time.time() < deadline:
            try:
                self.call("GET", "/api/health")
                break
            except OSError:
                time.sleep(0.2)
        else:
            raise TestAssertionFailed("the grr server did not answer within 30 s")
        login = self.call("POST", "/api/auth/login", {"use_env": True}, token="")
        self.token = login["token"]

    def call(self, method, path, body=None, token=None):
        req = urllib.request.Request(self.url + path, method=method, data=json.dumps(body).encode() if body is not None else None,
                                     headers={"Authorization": "Bearer " + (self.token if token is None else token), "Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=30) as res:
            return json.load(res)

    def stop(self):
        if self.proc is not None:
            try:
                os.killpg(self.proc.pid, 15)
                self.proc.wait(10)
            except (OSError, subprocess.TimeoutExpired):
                try:
                    os.killpg(self.proc.pid, 9)
                except OSError:
                    pass
        shutil.rmtree(self.home, ignore_errors=True)


class GrrTuiCase(TestCase):
    mode = "pty"
    timeout = 240.0
    extra_args = ("--tui",)

    async def setup(self):
        if not (GRR_DIR / "grr").is_dir():
            raise TestAssertionFailed("SKIPPED: no grr checkout at %s (set GRR_DIR)" % GRR_DIR)
        if not os.environ.get("DEVPLACE_API_KEY"):
            raise TestAssertionFailed("SKIPPED: DEVPLACE_API_KEY is not set")
        self.server = GrrServer()
        self.server.start()
        self.home = tempfile.mkdtemp(prefix="tai-it-")
        with open(os.path.join(self.home, "grr.json"), "w") as fh:
            json.dump({"url": self.server.url, "accounts": [{"token": self.server.token, "username": "it", "email": ""}], "active": self.server.token}, fh)
        os.chmod(os.path.join(self.home, "grr.json"), 0o600)
        self.env_overrides.update({"TAI_HOME": self.home, "TAI_PASSPHRASE": "", "GRR_URL": self.server.url, "TERM": "xterm-256color"})

    async def teardown(self):
        self.server.stop()
        shutil.rmtree(self.home, ignore_errors=True)

    async def start(self, rows=40, cols=130):
        """The real binary on a pty of a known size (set before the app reads it)."""
        import fcntl
        import struct
        import termios
        proc = await self.spawn(mode="pty")
        fcntl.ioctl(proc._master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
        return proc

    async def tui(self):
        proc = await self.start()
        await proc.wait_for_pattern(re.compile(r"Automations"), timeout=40)
        return proc

    def type(self, proc, text):
        os.write(proc._master_fd, text.encode())

    async def settle(self, proc, seconds=1.0):
        end = time.monotonic() + seconds
        while time.monotonic() < end:
            try:
                chunk = await asyncio.wait_for(proc._read_chunk(), timeout=0.2)
            except asyncio.TimeoutError:
                continue
            if chunk:
                proc._record("stdout", chunk.decode("utf-8", "replace"))

    def screen_text(self, proc, rows=40, cols=130):
        """What the terminal shows now: the bytes so far replayed through tai's own VT emulator."""
        sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
        import tai
        raw = "".join(l.text for l in proc.transcript if l.stream == "stdout")
        vt = tai.TuiVtScreen(rows, cols)
        vt.feed(raw.encode("utf-8", "replace"))
        lines = ["".join((c[0] if c else " ") for c in vt.grid[r]).rstrip() for r in range(rows)]
        if os.environ.get("GRR_IT_DEBUG"):
            print("\n".join(l for l in lines if l), flush=True)
            print("-" * 40, flush=True)
        return "\n".join(lines)

    def bot_chat(self, name):
        for bot in self.server.call("GET", "/api/bots")["bots"]:
            if bot["name"] == name:
                return bot, bot["chat"]["id"]
        raise TestAssertionFailed("no bot named %s" % name)


class DrawerAndNavigationTest(GrrTuiCase):
    """Deterministic: no model. The drawer lists the seed bots, Enter opens a
    chat, Esc goes back, the settings route opens by its shortcut and a switch
    flipped with Enter is saved by the server."""

    name = "grr-tui-drawer-and-settings"

    async def run(self):
        proc = await self.tui()
        await self.settle(proc, 1.5)
        text_before = "\n".join(l.text for l in proc.transcript if l.stream == "stdout")
        # the drawer names the three seed bots
        seen = strip_ansi(text_before)
        shown = {name: name in seen for name in ("Grr", "Smile", "Tai")}
        self.type(proc, "\x07")  # Ctrl+G: settings
        await proc.wait_for_pattern(re.compile(r"Auto-review"), timeout=15)
        # account, usage, connect apps, then the Auto-review switch: three steps down, then flip it
        before = self.server.call("GET", "/api/me")["settings"]["auto_review"]
        for _ in range(3):
            self.type(proc, "\x1b[B")
            await self.settle(proc, 0.3)
        self.type(proc, "\r")
        await self.settle(proc, 1.5)
        after = self.server.call("GET", "/api/me")["settings"]["auto_review"]
        self.type(proc, "\x03")
        return {"shown": shown, "before": before, "after": after}

    async def verify(self, result):
        missing = [name for name, ok in result["shown"].items() if not ok]
        if missing:
            raise TestAssertionFailed("the drawer did not show %s" % ", ".join(missing))
        flipped = [k for k in ("auto_review",) if result["before"] != result["after"]]
        if not flipped:
            raise TestAssertionFailed("the switch pressed with Enter did not change the server's setting (before=%r after=%r)" % (result["before"], result["after"]))


class ChatRoundTripTest(GrrTuiCase):
    """LLM-driven (a free, rotating backend: the exact model differs per run, so
    the reply text is not asserted, only that a real reply exists). Types a
    message to a bot in the TUI and checks, in the server's own message list,
    that the user message was stored, a bot reply followed, and that the reply
    reached the screen."""

    name = "grr-tui-chat-round-trip"

    async def run(self):
        bot, chat_id = self.bot_chat("Tai")
        proc = await self.tui()
        # open the Tai chat from the drawer by clicking its row, then type a message
        self.type(proc, "\x1b[B" * 5)
        await self.settle(proc, 0.6)
        self.type(proc, "\r")
        await proc.wait_for_pattern(re.compile(r"Ask Tai"), timeout=20)
        token = "pineapple%d" % (int(time.time()) % 1000)
        self.type(proc, "Reply with just the word %s and nothing else" % token)
        await self.settle(proc, 0.4)
        self.type(proc, "\r")
        reply = ""
        deadline = time.monotonic() + 150
        while time.monotonic() < deadline:
            await self.settle(proc, 2.0)
            messages = self.server.call("GET", "/api/chats/%s/messages" % chat_id)["messages"]
            replies = [m for m in messages if m["role"] == "assistant" and m["kind"] == "text"]
            if replies:
                reply = replies[-1]["content"]
                break
        await self.settle(proc, 2.0)
        screen = strip_ansi("\n".join(l.text for l in proc.transcript if l.stream == "stdout"))
        self.type(proc, "\x03")
        messages = self.server.call("GET", "/api/chats/%s/messages" % chat_id)["messages"]
        return {"token": token, "messages": messages, "reply": reply, "screen": screen}

    async def verify(self, result):
        mine = [m for m in result["messages"] if m["role"] == "user" and result["token"] in m["content"]]
        if not mine:
            raise TestAssertionFailed("the server never stored the message typed in the TUI")
        if not result["reply"]:
            raise TestAssertionFailed("no bot reply within 150 s (probabilistic: the free backend may have been down or slow)")
        if result["token"] not in result["screen"]:
            raise TestAssertionFailed("the typed message never appeared on the screen")
        if result["reply"][:12] not in result["screen"]:
            raise TestAssertionFailed("the reply %r exists on the server but never reached the screen" % result["reply"][:40])


class NewBotAndEditTest(GrrTuiCase):
    """Deterministic: no model. Creates a bot from the drawer's plus menu and
    renames it through the profile's editor dialog; the server must hold the new
    name."""

    name = "grr-tui-new-bot-and-rename"

    async def run(self):
        before = {b["id"] for b in self.server.call("GET", "/api/bots")["bots"]}
        proc = await self.tui()
        # the drawer: me, three nav rows, then the Bots header with its plus
        self.type(proc, "\x1b[B" * 4)
        await self.settle(proc, 0.5)
        self.screen_text(proc)
        self.type(proc, "\r")  # opens the menu
        await self.settle(proc, 0.8)
        self.screen_text(proc)
        self.type(proc, "\r")  # New Bot
        await proc.wait_for_pattern(re.compile(r"Documents"), timeout=20)
        await self.settle(proc, 1.0)
        # a new bot opens its profile with the name editor already open
        self.type(proc, "\x7f" * 20 + "Rocket")
        await self.settle(proc, 0.4)
        self.type(proc, "\r")
        await self.settle(proc, 2.0)
        self.type(proc, "\x03")
        after = self.server.call("GET", "/api/bots")["bots"]
        return {"new": [b for b in after if b["id"] not in before]}

    async def verify(self, result):
        if len(result["new"]) != 1:
            raise TestAssertionFailed("expected exactly one new bot, found %d" % len(result["new"]))
        if result["new"][0]["name"] != "Rocket":
            raise TestAssertionFailed("the bot was not renamed in the profile editor: its name is %r" % result["new"][0]["name"])


class NarrowPhoneTest(GrrTuiCase):
    """Deterministic: a 40 column terminal shows the drawer alone, Enter opens a
    chat full width, Esc returns, and no output line is wider than the terminal."""

    name = "grr-tui-narrow-terminal"

    async def run(self):
        proc = await self.start(30, 40)
        await proc.wait_for_pattern(re.compile(r"Automations"), timeout=40)
        await self.settle(proc, 1.0)
        self.type(proc, "\x1b[B" * 5)
        await self.settle(proc, 0.5)
        self.type(proc, "\r")
        await proc.wait_for_pattern(re.compile(r"Ask "), timeout=20)
        await self.settle(proc, 0.5)
        self.type(proc, "\x1b")
        await self.settle(proc, 1.0)
        screen = strip_ansi("\n".join(l.text for l in proc.transcript if l.stream == "stdout"))
        self.type(proc, "\x03")
        return {"screen": screen}

    async def verify(self, result):
        if "Ask " not in result["screen"]:
            raise TestAssertionFailed("the chat did not open on a 40 column terminal")
        if result["screen"].count("Automations") < 2:
            raise TestAssertionFailed("Esc did not bring the drawer back")


ALL_CASES = [DrawerAndNavigationTest(), NewBotAndEditTest(), NarrowPhoneTest(), ChatRoundTripTest()]

if __name__ == "__main__":
    sys.exit(main_cli(ALL_CASES, default_log_dir=str(Path(__file__).resolve().parent / "logs")))