STT-Assistant/tests/test_tui_client.py

"""Tests for the TUI control client against a real control server."""

import os
import shutil
import tempfile
import time

import pytest

from stt_daemon.assistant.control import ControlServer
from stt_daemon.assistant.service import AssistantService
from stt_daemon.config import Config
from stt_daemon.tui.client import ControlClient


class FakeSubscriber:
    def events(self):
        return iter(())

    def close(self):
        pass


@pytest.fixture
def server():
    directory = tempfile.mkdtemp(prefix="sttt-")
    path = os.path.join(directory, "c.sock")
    cfg = Config()
    cfg.assistant.control_socket = path

    service = AssistantService(cfg, FakeSubscriber())
    control = ControlServer(cfg, service, path=path)
    service.set_sinks(emit=control.broadcast, relay=control.relay)
    control.start()

    yield path, service, control

    control.close()
    service.stop()
    shutil.rmtree(directory, ignore_errors=True)


def _wait(predicate, timeout=5.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        if predicate():
            return True
        time.sleep(0.01)
    return predicate()


def test_client_subscribes_and_receives_status(server):
    path, _service, _control = server
    received: list[dict] = []

    client = ControlClient(path)
    client.connect()
    client.subscribe(viewer=True)
    client.start(received.append)

    assert _wait(lambda: any(e["type"] == "status" for e in received))
    client.close()


def test_client_command_start_flips_mode_and_broadcasts(server):
    path, service, _control = server
    received: list[dict] = []

    client = ControlClient(path)
    client.connect()
    client.subscribe(viewer=True)
    client.start(received.append)
    assert _wait(lambda: any(e["type"] == "status" for e in received))

    client.send("command_start")

    assert _wait(lambda: service.modes.mode.value == "command_armed")
    assert _wait(
        lambda: any(e["type"] == "mode" and e["mode"] == "command_armed" for e in received)
    )
    client.close()


def test_client_receives_relayed_transcript(server):
    path, service, _control = server
    received: list[dict] = []

    client = ControlClient(path)
    client.connect()
    client.subscribe(viewer=True)
    client.start(received.append)
    assert _wait(lambda: any(e["type"] == "status" for e in received))

    service.handle_event(
        {"v": 1, "type": "transcript", "seq": 1, "ts": time.time(), "text": "hello", "lang": "en"}
    )

    assert _wait(lambda: any(e["type"] == "transcript" and e["text"] == "hello" for e in received))
    client.close()


def test_client_reconnects_after_server_restart(server):
    path, service, control = server
    received: list[dict] = []
    states: list[bool] = []

    client = ControlClient(path)
    client.connect()
    client.subscribe(viewer=True)
    client.start(received.append, on_state=states.append)
    assert _wait(lambda: any(e["type"] == "status" for e in received))

    control.close()
    assert _wait(lambda: False in states)

    cfg = Config()
    cfg.assistant.control_socket = path
    control2 = ControlServer(cfg, service, path=path)
    service.set_sinks(emit=control2.broadcast, relay=control2.relay)
    control2.start()
    try:
        assert _wait(lambda: states[-1] is True if states else False)
        received.clear()
        assert _wait(lambda: any(e["type"] == "status" for e in received))
    finally:
        client.close()
        control2.close()