STT-Assistant/tests/test_daemon_events.py

"""Tests for the daemon's publish-only event stream (Phase 1)."""

import logging
import os
import shutil
import socket
import tempfile
import threading
import time

import numpy as np
import pytest

from stt_daemon import events
from stt_daemon.bus import EventSubscriber
from stt_daemon.config import Config
from stt_daemon.daemon import Daemon
from stt_daemon.transcriber import TranscriptResult

FRAME_SAMPLES = 480
FRAME_BYTES = FRAME_SAMPLES * 2


@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 _frame(value=0):
    return np.full(FRAME_SAMPLES, value, dtype=np.int16).tobytes()


class FakeVad:
    def __init__(self, decisions):
        self._decisions = list(decisions)
        self._index = 0

    def is_speech(self, frame, sample_rate):
        if self._index < len(self._decisions):
            decision = self._decisions[self._index]
            self._index += 1
            return decision
        return False


class FakeCapture:
    def __init__(self, frames, delay=0.0, stall_after=None, healthy=True):
        self._frames = list(frames)
        self._delay = delay
        self._stall_after = stall_after
        self._healthy = healthy
        self._closed = threading.Event()
        self.frame_bytes = FRAME_BYTES

    def frames(self):
        for index, frame in enumerate(self._frames):
            if self._closed.is_set():
                return
            if self._delay:
                time.sleep(self._delay)
            yield frame
            if self._stall_after is not None and index + 1 >= self._stall_after:
                self._closed.wait()
                return

    def healthy(self):
        return self._healthy

    def close(self):
        self._closed.set()


class FakeTranscriber:
    def __init__(self, text="hello", language="en", probability=0.99):
        self.text = text
        self.language = language
        self.probability = probability

    def transcribe(self, audio):
        return TranscriptResult(
            language=self.language,
            language_probability=self.probability,
            text=self.text,
        )

    def close(self):
        pass


class FakeSink:
    def __init__(self):
        self.emitted = []

    def emit(self, result, duration_s, rms=0.0, reason="unknown", audio=None):
        if not result.text:
            return None
        self.emitted.append(result.text)
        return None

    def close(self):
        pass


class _Collector:
    def __init__(self, subscriber):
        self.events = []
        self._thread = threading.Thread(target=self._run, args=(subscriber,), daemon=True)
        self._thread.start()

    def _run(self, subscriber):
        try:
            for event in subscriber.events():
                self.events.append(event)
        except Exception:  # noqa: BLE001 - collector must not crash the test thread
            pass

    def types(self):
        return [e["type"] for e in self.events]

    def of(self, type_):
        return [e for e in self.events if e["type"] == type_]

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


def _base_config(sock_path):
    cfg = Config()
    cfg.events.socket = sock_path
    cfg.events.heartbeat_ms = 50
    cfg.audio.frame_ms = 30
    cfg.vad.pre_roll_ms = 60
    cfg.vad.speech_start_ms = 30
    cfg.vad.silence_ms = 60
    cfg.vad.min_utterance_ms = 30
    cfg.vad.chunk_pause_ms = 300
    cfg.vad.soft_max_s = 10.0
    cfg.vad.max_utterance_s = 30.0
    cfg.vad.rms_floor = 0.0
    cfg.daemon.queue_size = 50
    cfg.output.stdout = False
    cfg.output.jsonl_dir = None
    cfg.output.save_audio = False
    return cfg


def _start_daemon(cfg, capture, vad, transcriber):
    sink = FakeSink()
    daemon = Daemon(
        cfg,
        log_level=logging.WARNING,
        capture_factory=lambda: capture,
        vad_factory=lambda _aggr: vad,
        transcriber_factory=lambda: transcriber,
        sink_factory=lambda: sink,
    )
    daemon._initialize()
    subscriber = EventSubscriber(cfg.events.socket, reconnect_delay=0.05)
    collector = _Collector(subscriber)
    daemon._start_workers()
    return daemon, collector, subscriber, sink


def test_transcript_and_segment_done_share_seg_id(sock_path):
    cfg = _base_config(sock_path)
    frames = [_frame() for _ in range(5)]
    vad = FakeVad([False, True, True, False, False])
    transcriber = FakeTranscriber(text="hello")
    daemon, collector, subscriber, _ = _start_daemon(cfg, FakeCapture(frames), vad, transcriber)

    collector.wait_for(lambda e: any(x["type"] == events.SEGMENT_DONE for x in e))

    starts = collector.of(events.SPEECH_START)
    stops = collector.of(events.SPEECH_STOP)
    transcripts = collector.of(events.TRANSCRIPT)
    dones = collector.of(events.SEGMENT_DONE)

    assert len(starts) == 1
    assert "seg_id" not in starts[0]  # speech_start carries no seg_id

    assert len(stops) == 1
    assert len(transcripts) == 1
    assert len(dones) == 1
    assert stops[0]["seg_id"] == transcripts[0]["seg_id"] == dones[0]["seg_id"] == 1
    assert transcripts[0]["text"] == "hello"
    assert dones[0]["emitted"] is True
    assert dones[0]["end_reason"] == "end_of_turn"

    subscriber.close()
    daemon.stop()


def test_gated_segment_reports_emitted_false(sock_path):
    cfg = _base_config(sock_path)
    frames = [_frame() for _ in range(5)]
    vad = FakeVad([False, True, True, False, False])
    transcriber = FakeTranscriber(text="")  # gated / no text
    daemon, collector, subscriber, _ = _start_daemon(cfg, FakeCapture(frames), vad, transcriber)

    collector.wait_for(lambda e: any(x["type"] == events.SEGMENT_DONE for x in e))

    assert collector.of(events.TRANSCRIPT) == []
    dones = collector.of(events.SEGMENT_DONE)
    assert len(dones) == 1
    assert dones[0]["emitted"] is False
    assert len(collector.of(events.SPEECH_STOP)) == 1

    subscriber.close()
    daemon.stop()


def test_chunked_utterance_yields_one_pair_per_chunk(sock_path):
    cfg = _base_config(sock_path)
    cfg.vad.soft_max_s = 0.3  # 10 frames
    cfg.vad.chunk_pause_ms = 60  # 2 frames
    cfg.vad.silence_ms = 600  # 20 frames (never reached here)
    frames = [_frame() for _ in range(47)]
    decisions = [False] + [True] * 12 + [False] * 2 + [True] * 12 + [False] * 20
    transient = FakeTranscriber(text="chunk")
    daemon, collector, subscriber, _ = _start_daemon(
        cfg, FakeCapture(frames), FakeVad(decisions), transient
    )

    collector.wait_for(
        lambda e: len([x for x in e if x["type"] == events.SEGMENT_DONE]) >= 2
    )

    stops = collector.of(events.SPEECH_STOP)
    dones = collector.of(events.SEGMENT_DONE)
    assert len(stops) == len(dones) >= 2

    stop_ids = [s["seg_id"] for s in stops]
    done_ids = [d["seg_id"] for d in dones]
    assert stop_ids == sorted(stop_ids)
    assert len(set(stop_ids)) == len(stop_ids)  # unique per chunk
    assert set(stop_ids) == set(done_ids)
    assert all(d["end_reason"] == "pause" for d in dones)

    subscriber.close()
    daemon.stop()


def test_heartbeats_advance_and_carry_health(sock_path):
    cfg = _base_config(sock_path)
    cfg.events.heartbeat_ms = 40
    frames = [_frame() for _ in range(40)]
    daemon, collector, subscriber, _ = _start_daemon(
        cfg, FakeCapture(frames, delay=0.02), FakeVad([]), FakeTranscriber()
    )

    collector.wait_for(lambda e: len([x for x in e if x["type"] == events.HEARTBEAT]) >= 3)
    heartbeats = collector.of(events.HEARTBEAT)

    assert len(heartbeats) >= 3
    seqs = [h["seq"] for h in heartbeats]
    assert seqs == sorted(seqs)
    assert len(set(seqs)) == len(seqs)

    last = heartbeats[-1]
    assert last["model_loaded"] is True
    assert last["device_ok"] is True
    assert last["capture"] == cfg.audio.device
    for field in ("in_flight", "segments_done", "worker_stall_s", "queue_depth",
                  "frames_since_last", "rms", "uptime_s", "last_transcript_age_s"):
        assert field in last

    subscriber.close()
    daemon.stop()


def test_stalled_capture_stops_heartbeats(sock_path):
    cfg = _base_config(sock_path)
    cfg.events.heartbeat_ms = 60
    frames = [_frame() for _ in range(30)]
    capture = FakeCapture(frames, delay=0.01, stall_after=10)
    daemon, collector, subscriber, _ = _start_daemon(
        cfg, capture, FakeVad([]), FakeTranscriber()
    )

    assert collector.wait_for(
        lambda e: any(x["type"] == events.HEARTBEAT for x in e), timeout=3.0
    )
    time.sleep(0.2)  # let it stall
    before = len(collector.of(events.HEARTBEAT))

    time.sleep(0.45)  # ~7 heartbeat intervals at 60ms
    after = len(collector.of(events.HEARTBEAT))
    assert after == before, f"heartbeats continued during a stall ({before} -> {after})"

    subscriber.close()
    daemon.stop()


def test_discarded_segment_still_emits_speech_stop(sock_path):
    """A turn that is discarded (blip / too quiet) must still emit speech_stop.

    Without it, consumers that track speech_open from speech_start/speech_stop
    latch true forever, which blocks readiness and chat dispatch.
    """
    cfg = _base_config(sock_path)
    cfg.vad.min_utterance_ms = 300  # a 1-frame blip is discarded
    frames = [_frame() for _ in range(4)]
    vad = FakeVad([True, False, False, False])  # confirm, then silence -> discard
    daemon, collector, subscriber, _ = _start_daemon(
        cfg, FakeCapture(frames), vad, FakeTranscriber()
    )

    collector.wait_for(lambda e: any(x["type"] == events.SPEECH_STOP for x in e))
    time.sleep(0.2)

    starts = collector.of(events.SPEECH_START)
    stops = collector.of(events.SPEECH_STOP)

    assert len(starts) == 1
    assert len(stops) == 1
    assert stops[0]["seg_id"] is None  # discarded: no segment assigned
    assert stops[0]["end_reason"] == "discarded"
    assert collector.of(events.SEGMENT_DONE) == []
    assert collector.of(events.TRANSCRIPT) == []

    subscriber.close()
    daemon.stop()


def test_inbound_data_on_subscriber_socket_is_ignored(sock_path):
    cfg = _base_config(sock_path)
    cfg.events.heartbeat_ms = 40
    frames = [_frame() for _ in range(60)]
    daemon, collector, subscriber, _ = _start_daemon(
        cfg, FakeCapture(frames, delay=0.02), FakeVad([]), FakeTranscriber()
    )

    # The bus is publish-only: a client writing data must have no effect.
    raw = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    raw.connect(sock_path)
    raw.sendall(b'{"v":1,"type":"evil","seq":1,"ts":0}\n')
    time.sleep(0.1)

    assert collector.wait_for(
        lambda e: len([x for x in e if x["type"] == events.HEARTBEAT]) >= 2
    )
    assert all(x["type"] != "evil" for x in collector.events)

    raw.close()
    subscriber.close()
    daemon.stop()


def test_request_stop_shuts_down_run_loop(tmp_path):
    """request_stop() (the signal-handler path) must unwind run() cleanly."""
    sock_path = str(tmp_path / "events.sock")
    cfg = _base_config(sock_path)
    capture = FakeCapture([_frame(0)] * 5, stall_after=1)
    sink = FakeSink()
    daemon = Daemon(
        cfg,
        log_level=logging.WARNING,
        capture_factory=lambda: capture,
        vad_factory=lambda _aggr: FakeVad([]),
        transcriber_factory=lambda: FakeTranscriber(),
        sink_factory=lambda: sink,
    )

    thread = threading.Thread(target=daemon.run, daemon=True)
    thread.start()
    time.sleep(0.2)
    daemon.request_stop()
    thread.join(timeout=5)

    assert not thread.is_alive()
    assert capture._closed.is_set()