STT-Assistant/tests/test_assistant_integration.py

"""Integration: the assistant service consuming a real publish-only bus."""

import json
import os
import shutil
import signal
import sys
import tempfile
import threading
import time
from types import SimpleNamespace

import pytest

from stt_daemon import events
from stt_daemon.assistant.runner import build_triage
from stt_daemon.assistant.service import AssistantService
from stt_daemon.bus import EventPublisher, EventSubscriber
from stt_daemon.config import Config


@pytest.fixture
def sock_path():
    directory = tempfile.mkdtemp(prefix="sttb-")
    try:
        yield os.path.join(directory, "s.sock")
    finally:
        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 _heartbeat_fields():
    return {
        "in_flight": False,
        "segments_done": 0,
        "worker_stall_s": 0.0,
        "queue_depth": 0,
        "frames_since_last": 1,
        "rms": 0.0,
        "capture": "test",
        "device_ok": True,
        "model_loaded": True,
        "uptime_s": 1.0,
        "last_transcript_age_s": 0.0,
    }


def test_service_consumes_real_bus(sock_path):
    cfg = Config()
    publisher = EventPublisher(sock_path, ring_size=10)
    publisher.start()

    subscriber = EventSubscriber(sock_path, reconnect_delay=0.05)
    emitted: list[dict] = []
    service = AssistantService(
        cfg,
        subscriber,
        emit=lambda type_, **fields: emitted.append({"type": type_, **fields}),
    )
    service.start()

    try:
        publisher.emit(events.HEARTBEAT, **_heartbeat_fields())
        assert _wait(lambda: service.stt_ok() is True)

        assert service.arm() is True
        assert _wait(lambda: any(e.get("mode") == "command_armed" for e in emitted))

        publisher.emit(events.SPEECH_START)
        assert _wait(lambda: service.speech_open is True)

        publisher.emit(events.SPEECH_STOP, seg_id=5)
        assert _wait(lambda: service.awaited_seg_id == 5)
        assert service.ready() is False

        publisher.emit(events.SEGMENT_DONE, seg_id=5, emitted=True)
        assert _wait(lambda: service.ready() is True)

        publisher.emit(events.TRANSCRIPT, seg_id=5, text="hello", lang="en", lang_prob=0.9)
        assert _wait(lambda: service.command_text == "hello")

        assert service.execute() == "hello"
        assert service.finish() is True
    finally:
        service.stop()
        publisher.close()


def _enabled_triage_cfg(tmp_path):
    cfg = Config()
    cfg.audit.enabled = False
    cfg.llm.enabled = False
    cfg.triage.enabled = True
    path = os.path.join(str(tmp_path), "questions.json")
    with open(path, "w", encoding="utf-8") as handle:
        json.dump({"urgent": {"type": "noul", "instructions": "is it urgent"}}, handle)
    cfg.triage.questions_file = path
    return cfg


class _FakeChatClient:
    def __init__(self):
        self.calls = []

    def chat(self, messages, tools=None, max_tokens=None, timeout_s=None):
        self.calls.append(
            {
                "messages": messages,
                "tools": tools,
                "max_tokens": max_tokens,
                "timeout_s": timeout_s,
            }
        )
        return SimpleNamespace(content="EXTRACTED")


def test_build_triage_disabled_returns_nones():
    cfg = Config()
    assert cfg.triage.enabled is False
    assert build_triage(cfg, None) == (None, None, None)
    assert build_triage(cfg, _FakeChatClient()) == (None, None, None)
    assert "laya" not in sys.modules
    service = AssistantService(cfg, subscriber=object())
    assert service.status()["triage_state"] == "disabled"


def test_build_triage_enabled_constructs_without_client(tmp_path):
    cfg = _enabled_triage_cfg(tmp_path)
    before = len(threading.enumerate())
    analyzer, lane, extract_fn = build_triage(cfg, None)
    assert analyzer is not None
    assert lane is not None
    assert extract_fn is None
    assert analyzer._cfg is cfg.triage
    assert lane._cfg is cfg.triage
    assert analyzer.state() == ("loading", "")
    assert len(threading.enumerate()) == before
    assert "laya" not in sys.modules


def test_build_triage_extract_fn_calls_client(tmp_path):
    cfg = _enabled_triage_cfg(tmp_path)
    client = _FakeChatClient()
    _, _, extract_fn = build_triage(cfg, client)
    assert extract_fn is not None
    assert extract_fn("PROMPT") == "EXTRACTED"
    assert len(client.calls) == 1
    call = client.calls[0]
    assert call["messages"] == [
        {"role": "system", "content": cfg.llm.triage_system_prompt},
        {"role": "user", "content": "PROMPT"},
    ]
    assert call["max_tokens"] == cfg.triage.extract_max_tokens
    assert call["timeout_s"] == cfg.triage.extract_timeout_s


def test_build_triage_on_state_refreshes_analyzer(tmp_path):
    from stt_daemon.assistant.laya_lane import LAYA_MISSING_HINT

    cfg = _enabled_triage_cfg(tmp_path)
    sentinel = object()
    analyzer, lane, _ = build_triage(cfg, None, on_result=sentinel)
    assert analyzer is not None
    assert lane is not None
    assert lane._on_result is sentinel
    assert callable(lane._on_state)
    lane._on_state("ok", "")
    assert analyzer.state() == ("ok", "")
    lane._on_state("laya_missing", LAYA_MISSING_HINT)
    assert analyzer.state() == ("laya_missing", LAYA_MISSING_HINT)


def _run_runner_main(monkeypatch, cfg):
    from stt_daemon.assistant import runner as runner_mod
    import stt_daemon.assistant.laya_lane as lane_mod

    calls = []
    made_lanes = []
    made_services = []

    class FakeLane:
        def __init__(self, triage_cfg, agent_factory=None, on_result=None, on_state=None):
            made_lanes.append(self)
            self.cfg = triage_cfg
            self.on_state = on_state
            self.on_result = on_result
            calls.append("lane.construct")

        def start(self):
            calls.append("lane.start")

        def stop(self):
            calls.append("lane.stop")

    class FakeSubscriber:
        def close(self):
            calls.append("subscriber.close")

    class FakeService:
        def __init__(self, cfg, subscriber, **kwargs):
            made_services.append(self)
            self.kwargs = kwargs
            calls.append("service.construct")

        def set_sinks(self, **kwargs):
            pass

        def chat_notify(self, *args, **kwargs):
            return None

        def _on_triage_result(self, generation, payload):
            calls.append("triage.result")

        def start(self):
            calls.append("service.start")

        def stop(self):
            calls.append("service.stop")

    class FakeControl:
        def __init__(self, cfg, service):
            self.service = service
            calls.append("control.construct")

        path = "/tmp/fake.sock"

        def start(self):
            calls.append("control.start")

        def broadcast(self, *args, **kwargs):
            pass

        def relay(self, event):
            pass

        def close(self):
            calls.append("control.close")

    class FakeEvent:
        def set(self):
            pass

        def wait(self, timeout=None):
            return True

    monkeypatch.setattr(lane_mod, "LayaLane", FakeLane)
    monkeypatch.setattr(runner_mod, "EventSubscriber", lambda *a, **k: FakeSubscriber())
    monkeypatch.setattr(runner_mod, "ControlServer", FakeControl)
    monkeypatch.setattr(runner_mod, "AssistantService", FakeService)
    monkeypatch.setattr(runner_mod, "load_config", lambda path=None: cfg)
    monkeypatch.setattr(threading, "Event", FakeEvent)
    monkeypatch.setattr(signal, "signal", lambda *a, **k: None)
    monkeypatch.setattr(sys, "argv", ["stt-assistant"])
    rc = runner_mod.main()
    return rc, calls, made_lanes, made_services


def test_runner_main_shutdown_order_lane_before_control(tmp_path, monkeypatch):
    cfg = _enabled_triage_cfg(tmp_path)
    rc, calls, made_lanes, made_services = _run_runner_main(monkeypatch, cfg)
    assert rc == 0
    assert len(made_lanes) == 1
    assert made_lanes[0].cfg is cfg.triage
    assert made_services[0].kwargs["laya_lane"] is made_lanes[0]
    assert made_services[0].kwargs["triage"] is not None
    assert made_services[0].kwargs["triage_extract_fn"] is None
    assert calls.index("control.start") < calls.index("service.start")
    assert calls.index("service.start") < calls.index("lane.start")
    assert calls.index("service.stop") < calls.index("lane.stop")
    assert calls.index("lane.stop") < calls.index("control.close")
    made_lanes[0].on_result(3, {"answers": {}})
    assert "triage.result" in calls
    assert "laya" not in sys.modules


def test_runner_main_disabled_builds_no_lane(monkeypatch):
    cfg = Config()
    cfg.audit.enabled = False
    cfg.llm.enabled = False
    rc, calls, made_lanes, made_services = _run_runner_main(monkeypatch, cfg)
    assert rc == 0
    assert made_lanes == []
    assert "lane.construct" not in calls
    assert "lane.start" not in calls
    assert "lane.stop" not in calls
    assert made_services[0].kwargs["triage"] is None
    assert made_services[0].kwargs["laya_lane"] is None
    assert made_services[0].kwargs["triage_extract_fn"] is None
    assert not [t for t in threading.enumerate() if t.name.startswith("laya-")]