prain: terminal UI

2026-10-27 13:52:16 +01:00
parent 041398d849
commit 0531aaed44
2 changed files with 185 additions and 0 deletions
@@ -0,0 +1,70 @@
+import fcntl
+import os
+import pty
+import struct
+import subprocess
+import sys
+import termios
+import threading
+import unittest
+
+from prain.tui import Chat
+from tests.helpers import StubVllm, wait_for
+
+
+class Tui(unittest.TestCase):
+    def setUp(self):
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+
+    def test_typing_warms_up_and_enter_streams_the_answer(self):
+        chat = Chat(self.stub.url, "stub-model")
+        chat.type("Why is the sky")
+        wait_for(lambda: self.stub.bodies and not chat.warmer.busy)
+        self.assertEqual(self.stub.bodies[0]["priority"], 1)
+        chat.send()
+        chat.type("next question")  # an answer is streaming: no warm-up in "idle" mode
+        wait_for(lambda: chat.poll() or chat.reply is None)
+        self.assertEqual(chat.messages, [{"role": "user", "content": "Why is the sky"},
+                                         {"role": "assistant", "content": "Hello!"}])
+        self.assertEqual(len(self.stub.bodies), 2)
+        self.assertIn("tok/s", chat.status)
+
+    def test_warmup_off_and_cancel_on_send(self):
+        off = Chat(self.stub.url, "stub-model", warmup="off")
+        off.type("hello")
+        self.assertFalse(off.warmer.busy)
+        self.stub.delay = 3.0
+        chat = Chat(self.stub.url, "stub-model")
+        chat.type("a slow warm-up")
+        wait_for(lambda: self.stub.bodies)
+        chat.send()  # aborts the warm-up instead of queueing behind it
+        wait_for(lambda: not chat.warmer.busy, timeout=2.0)
+
+    def test_curses_screen_runs_in_a_real_terminal(self):
+        master, slave = pty.openpty()
+        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 80, 0, 0))
+        tui = subprocess.Popen([sys.executable, "-m", "prain.cli", "tui", "--url", self.stub.url],
+                               stdin=slave, stdout=slave, stderr=slave,
+                               env={**os.environ, "TERM": "xterm"}, start_new_session=True)
+        os.close(slave)
+        screen = bytearray()
+        threading.Thread(target=self._drain, args=(master, screen), daemon=True).start()
+        wait_for(lambda: b"stub-model" in screen)  # status line drawn
+        os.write(master, b"hi\r")
+        wait_for(lambda: any(body.get("stream") and "priority" not in body
+                             for body in self.stub.bodies))
+        wait_for(lambda: b"tok/s" in screen)
+        os.write(master, b"\x1b")
+        self.assertEqual(tui.wait(timeout=5), 0)
+
+    @staticmethod
+    def _drain(fd, screen):
+        while True:
+            try:
+                data = os.read(fd, 4096)
+            except OSError:
+                return
+            if not data:
+                return
+            screen.extend(data)