STT-Assistant/tests/test_assistant_service.py

"""Tests for the assistant service (Phase 2) using synthetic event streams."""

import itertools
import threading
import time

from stt_daemon import events
from stt_daemon.assistant.service import AssistantService
from stt_daemon.config import Config

_sequence = itertools.count(1)


class Clock:
    def __init__(self, start=1000.0):
        self.t = float(start)

    def __call__(self):
        return self.t

    def advance(self, dt):
        self.t += float(dt)


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

    def close(self):
        pass


def _service():
    cfg = Config()
    clock = Clock()
    emitted: list[dict] = []

    def emit(type_, **fields):
        emitted.append({"type": type_, **fields})

    service = AssistantService(cfg, FakeSubscriber(), emit=emit, time_fn=clock)
    return service, emitted, clock


def _ev(type_, clock, **fields):
    event = {"v": 1, "type": type_, "seq": next(_sequence), "ts": clock()}
    event.update(fields)
    return event


def _transcript(clock, text, ts=None):
    return {
        "v": 1,
        "type": events.TRANSCRIPT,
        "seq": next(_sequence),
        "ts": clock() if ts is None else ts,
        "text": text,
    }


def test_ready_requires_matching_segment_done():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.ready() is False

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    assert service.ready() is False  # awaiting segment 1

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))
    assert service.ready() is True


def test_stale_segment_done_does_not_satisfy_later_turn():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SPEECH_START, clock))
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=2))

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))
    assert service.ready() is False  # still awaiting segment 2

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=2, emitted=True))
    assert service.ready() is True


def test_gated_segment_done_clears_readiness():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=3))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=3, emitted=False))

    assert service.ready() is True


def test_speech_stop_without_seg_id_clears_readiness():
    """A discarded turn closes with speech_stop but no seg_id; readiness must recover."""
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.ready() is False

    service.handle_event(_ev(events.SPEECH_STOP, clock))  # discarded: no seg_id
    assert service.speech_open is False
    assert service.awaited_seg_id is None
    assert service.ready() is True


def test_status_resets_stale_speech_state_on_reconnect():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.speech_open is True

    service.handle_event(_ev(events.STATUS, clock, capture="x", segments_done=0))
    assert service.speech_open is False
    assert service.awaited_seg_id is None
    assert service.ready() is True


def test_command_buffer_uses_arm_timestamp():
    service, _, clock = _service()
    clock.advance(10)
    assert service.arm() is True
    arm_ts = service.arm_ts
    assert arm_ts is not None

    # Replayed transcript from before arming is ignored.
    service.handle_event(_transcript(clock, "old", ts=arm_ts - 5))
    assert service.command_text == ""

    service.handle_event(_transcript(clock, "what time is it"))
    assert service.command_text == "what time is it"


def test_mutual_exclusion_between_command_and_chat():
    service, _, _ = _service()

    assert service.toggle_chat() is True
    assert service.arm() is False  # cannot arm in chat

    assert service.toggle_chat() is True  # back to idle
    assert service.arm() is True
    assert service.toggle_chat() is False  # cannot chat while armed


def test_arm_timeout_disarms_and_emits_error():
    service, emitted, clock = _service()
    service.arm()

    clock.advance(service.cfg.assistant.command_arm_timeout_s + 0.1)
    service.tick()

    assert service.modes.mode.value == "idle"
    assert service.command_text == ""
    assert any(e["type"] == "error" and e["code"] == "arm_timeout" for e in emitted)


def test_long_speech_does_not_disarm():
    service, emitted, clock = _service()
    service.arm()
    service.handle_event(_ev(events.SPEECH_START, clock))

    clock.advance(service.cfg.assistant.command_arm_timeout_s + 30.0)
    service.tick()

    assert service.modes.mode.value == "command_armed"
    assert not any(e["type"] == "error" for e in emitted)


def test_captured_command_survives_arm_timeout():
    service, emitted, clock = _service()
    service.arm()
    service.handle_event(_transcript(clock, "what time is it"))
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))

    clock.advance(service.cfg.assistant.command_arm_timeout_s + 30.0)
    service.tick()

    assert service.modes.mode.value == "command_armed"
    assert service.command_text == "what time is it"
    assert not any(e["type"] == "error" for e in emitted)


def test_stt_stall_cancels_armed_command():
    service, emitted, clock = _service()
    service.handle_event(_ev(events.HEARTBEAT, clock, worker_stall_s=0.0))
    service.arm()

    clock.advance(service.cfg.assistant.stall_timeout_ms / 1000.0 + 0.5)
    service.tick()

    assert service.modes.mode.value == "idle"
    assert any(e["type"] == "error" and e["code"] == "stt_stall" for e in emitted)


def test_worker_stall_marks_stt_degraded():
    service, _, clock = _service()
    long_stall = service.cfg.assistant.worker_stall_ms / 1000.0 + 1.0

    service.handle_event(_ev(events.HEARTBEAT, clock, worker_stall_s=long_stall))
    assert service.stt_ok() is False
    assert service.ready() is False


def test_chat_idle_timeout_exits_chat():
    service, _, clock = _service()
    assert service.toggle_chat() is True

    clock.advance(service.cfg.assistant.chat_idle_timeout_s + 0.1)
    service.tick()

    assert service.modes.mode.value == "idle"


def test_chat_exit_phrase_exits_chat():
    service, _, clock = _service()
    assert service.toggle_chat() is True

    service.handle_event(_transcript(clock, "That's all."))

    assert service.modes.mode.value == "idle"
    assert service.chat_history == []


def test_chat_accumulates_turns():
    service, _, clock = _service()
    service.toggle_chat()

    for text in ("hello there", "how are you"):
        service.handle_event(_transcript(clock, text))

    assert service.chat_history == ["hello there", "how are you"]


def test_execute_requires_readiness():
    service, _, clock = _service()
    service.arm()
    service.handle_event(_transcript(clock, "do it"))

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.execute() is None  # speech still open
    assert service.modes.mode.value == "command_armed"

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    assert service.execute() is None  # awaiting segment 1

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))
    assert service.execute() == "do it"
    assert service.modes.mode.value == "executing"

    assert service.finish() is True
    assert service.modes.mode.value == "idle"


def test_turn_ended_requires_quiet_gap():
    service, _, clock = _service()
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))

    assert service.ready() is True
    assert service.turn_ended() is False

    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    assert service.turn_ended() is True


def test_emits_mode_and_command_text():
    service, emitted, clock = _service()
    service.arm()
    assert {"type": "mode", "mode": "command_armed"} in emitted

    service.handle_event(_transcript(clock, "hi"))
    assert any(e["type"] == "command_text" and e["text"] == "hi" for e in emitted)


def test_cancel_clears_command_buffer():
    service, _, clock = _service()
    service.arm()
    service.handle_event(_transcript(clock, "hi"))
    assert service.cancel() is True
    assert service.modes.mode.value == "idle"
    assert service.command_text == ""
    assert service.arm_ts is None


def test_pending_chat_is_not_lost_across_threads():
    """Transcripts appended on the reader thread while the ticker dispatches
    must all reach the model: the join-and-clear in _dispatch_chat is atomic."""
    cfg = Config()
    cfg.assistant.turn_gap_s = 0.01
    dispatched: list[str] = []

    def emit(type_, **fields):
        if type_ == "chat_message" and fields.get("role") == "user":
            dispatched.append(fields["text"])

    service = AssistantService(cfg, FakeSubscriber(), emit=emit, chat_fn=lambda text: "ok")
    service.toggle_chat()

    total = 200
    producer = threading.Thread(
        target=lambda: [
            service.handle_event(
                {
                    "v": 1,
                    "type": events.TRANSCRIPT,
                    "seq": i,
                    "ts": time.time(),
                    "text": f"w{i}",
                }
            )
            for i in range(total)
        ],
        daemon=True,
    )
    producer.start()

    deadline = time.time() + 10.0
    while time.time() < deadline:
        service.tick()
        words = sum(len(text.split()) for text in dispatched)
        if not producer.is_alive() and words == total:
            break
        time.sleep(0.002)
    producer.join(timeout=5)

    assert sum(len(text.split()) for text in dispatched) == total