71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
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)
|