STT-Assistant/tests/test_daemon.py

"""Tests for the daemon's worker processing and transcription backlog warning."""

import json
import logging

import numpy as np
import pytest

from stt_daemon import events
from stt_daemon import transcriber as transcriber_mod
from stt_daemon.config import Config
from stt_daemon.daemon import Daemon
from stt_daemon.output import TranscriptSink
from stt_daemon.segmenter import Utterance
from stt_daemon.transcriber import TranscriptResult


def _daemon(threshold=2, queue_size=10):
    cfg = Config()
    cfg.daemon.backlog_warn_at = threshold
    cfg.daemon.queue_size = queue_size
    return Daemon(cfg)


def test_queue_size_comes_from_config():
    daemon = _daemon(queue_size=7)
    assert daemon.segment_queue.maxsize == 7


def test_backlog_warning_is_edge_triggered(caplog):
    daemon = _daemon(threshold=2)

    with caplog.at_level(logging.WARNING, logger="stt_daemon"):
        daemon.segment_queue.put("a")
        daemon._check_backlog()
        assert "backlog" not in caplog.text.lower()

        daemon.segment_queue.put("b")
        daemon._check_backlog()
        assert caplog.text.lower().count("backlog") == 1

        daemon._check_backlog()
        assert caplog.text.lower().count("backlog") == 1

        daemon.segment_queue.get()
        daemon.segment_queue.get()
        daemon._check_backlog()
        assert daemon._backlog_warned is False

        daemon.segment_queue.put("c")
        daemon.segment_queue.put("d")
        daemon._check_backlog()
        assert caplog.text.lower().count("backlog") == 2


class _FakeWhisperModel:
    """Constructor-only fake; keeps faster-whisper from loading a real model."""

    def __init__(self, *args, **kwargs):
        self.init_kwargs = kwargs


class _FakeCppModel:
    """Constructor-only fake; keeps pywhispercpp from loading a real model."""

    def __init__(self, *args, **kwargs):
        self.init_kwargs = kwargs


def test_build_transcriber_defaults_to_faster_whisper(monkeypatch):
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _FakeWhisperModel)
    assert isinstance(_daemon()._build_transcriber(), transcriber_mod.Transcriber)


def test_build_transcriber_honors_pywhispercpp_backend(monkeypatch):
    cpp_model = pytest.importorskip("pywhispercpp.model")
    monkeypatch.setattr(cpp_model, "Model", _FakeCppModel)
    from stt_daemon.pywhispercpp_transcriber import PyWhisperCppTranscriber

    cfg = Config()
    cfg.whisper.backend = "pywhispercpp"
    assert isinstance(Daemon(cfg)._build_transcriber(), PyWhisperCppTranscriber)


# --- worker processing ------------------------------------------------------


class _FakeTranscriber:
    """Returns a canned result without loading a real model."""

    def __init__(self, result):
        self.result = result

    def transcribe(self, audio):
        return self.result

    def close(self):
        pass


class _RaisingSink:
    """Sink that fails on emit, simulating a non-JSON-serializable record."""

    def emit(self, *args, **kwargs):
        raise TypeError("Object of type float32 is not JSON serializable")

    def close(self):
        pass


def _utterance(seg_id=1):
    return Utterance(
        audio=np.zeros(1600, dtype=np.float32),
        duration_s=1.0,
        started_at=0.0,
        rms=0.05,
        reason="pause",
        seg_id=seg_id,
    )


def _transcript_result():
    return TranscriptResult(language="en", language_probability=0.9, text="hello world")


def _daemon_with_raising_sink():
    daemon = _daemon()
    daemon.transcriber = _FakeTranscriber(_transcript_result())
    daemon.sink = _RaisingSink()
    published = []
    daemon._publish = lambda type_, **fields: published.append((type_, fields))
    return daemon, published


def test_sink_failure_still_publishes_transcript(caplog):
    daemon, published = _daemon_with_raising_sink()

    with caplog.at_level(logging.ERROR, logger="stt_daemon"):
        daemon._process(_utterance())

    transcripts = [fields for type_, fields in published if type_ == events.TRANSCRIPT]
    assert len(transcripts) == 1
    assert transcripts[0]["text"] == "hello world"
    assert transcripts[0]["lang_prob"] == 0.9
    assert transcripts[0]["audio"] is None
    assert "Output error" in caplog.text
    sid = daemon._session_tracker.current_id
    assert transcripts[0]["session_id"] == sid
    assert sid
    assert transcripts[0]["session_dir"].endswith(sid)


def test_sink_failure_reports_segment_done_as_emitted():
    daemon, published = _daemon_with_raising_sink()

    daemon._process(_utterance())

    done = [fields for type_, fields in published if type_ == events.SEGMENT_DONE]
    assert len(done) == 1
    assert done[0]["emitted"] is True


# --- session identity -------------------------------------------------------


class _RecordingSink:
    """Minimal sink capturing the per-emit session_dir kwarg."""

    def __init__(self):
        self.session_dirs = []

    def emit(self, *args, **kwargs):
        self.session_dirs.append(kwargs.get("session_dir"))
        return None

    def close(self):
        pass


def _session_daemon(monkeypatch, clock):
    daemon = _daemon()
    daemon.transcriber = _FakeTranscriber(_transcript_result())
    sink = _RecordingSink()
    daemon.sink = sink
    published = []
    daemon._publish = lambda type_, **fields: published.append((type_, fields))
    monkeypatch.setattr("stt_daemon.daemon.time.time", lambda: clock["now"])
    return daemon, sink, published


def test_process_stamps_session_id_matching_jsonl(tmp_path):
    daemon = _daemon()
    daemon.transcriber = _FakeTranscriber(_transcript_result())
    daemon.cfg.output.stdout = False
    daemon.cfg.output.jsonl = True
    daemon.cfg.output.save_audio = False
    daemon.cfg.output.session_dir = str(tmp_path)
    daemon.sink = TranscriptSink(stdout=False, jsonl=True, save_audio=False, sample_rate=16000)
    published = []
    daemon._publish = lambda type_, **fields: published.append((type_, fields))

    daemon._process(_utterance(1))

    sid = daemon._session_tracker.current_id
    assert isinstance(sid, str) and sid
    transcripts = [f for t, f in published if t == events.TRANSCRIPT]
    assert len(transcripts) == 1
    assert transcripts[0]["session_id"] == sid
    assert transcripts[0]["session_dir"] == str(tmp_path / sid)
    sessions = [f for t, f in published if t == events.SESSION]
    assert len(sessions) == 1
    assert sessions[0]["session_id"] == sid

    record = json.loads((tmp_path / sid / "transcripts.jsonl").read_text())
    assert record["session_id"] == sid


def test_within_gap_reuses_one_session(monkeypatch):
    clock = {"now": 1000.0}
    daemon, sink, published = _session_daemon(monkeypatch, clock)

    daemon._process(_utterance(1))
    clock["now"] = 1100.0
    daemon._process(_utterance(2))

    sessions = [f for t, f in published if t == events.SESSION]
    assert len(sessions) == 1
    transcripts = [f for t, f in published if t == events.TRANSCRIPT]
    assert len(transcripts) == 2
    assert transcripts[0]["session_id"] == transcripts[1]["session_id"]
    assert transcripts[0]["session_id"] == sessions[0]["session_id"]
    assert len(sink.session_dirs) == 2
    assert all(
        d is not None and d.name == sessions[0]["session_id"] for d in sink.session_dirs
    )


def test_beyond_gap_mints_new_session(monkeypatch):
    clock = {"now": 1000.0}
    daemon, _sink, published = _session_daemon(monkeypatch, clock)

    daemon._process(_utterance(1))
    clock["now"] = 1100.0  # inside the 300s default gap
    daemon._process(_utterance(2))
    clock["now"] = 1500.0  # 400s > 300s: rollover
    daemon._process(_utterance(3))

    sessions = [f for t, f in published if t == events.SESSION]
    assert len(sessions) == 2
    assert sessions[0]["session_id"] != sessions[1]["session_id"]
    transcripts = [f for t, f in published if t == events.TRANSCRIPT]
    assert transcripts[0]["session_id"] == sessions[0]["session_id"]
    assert transcripts[1]["session_id"] == sessions[0]["session_id"]
    assert transcripts[2]["session_id"] == sessions[1]["session_id"]
    assert daemon._session_tracker.current_id == sessions[1]["session_id"]


def test_session_advances_with_file_writes_off(monkeypatch):
    clock = {"now": 1000.0}
    daemon, _sink, published = _session_daemon(monkeypatch, clock)
    daemon.cfg.output.jsonl = False
    daemon.cfg.output.save_audio = False

    daemon._process(_utterance(1))
    clock["now"] = 1500.0
    daemon._process(_utterance(2))

    sessions = [f for t, f in published if t == events.SESSION]
    assert len(sessions) == 2
    assert sessions[1]["session_id"] != sessions[0]["session_id"]