STT-Assistant/tests/test_daemon_speaker.py

"""Daemon speaker-identification lane: embedder seam, join-before-emit, inbox.

All hardware- and network-free: a fake embedder returns fixed vectors and the
store lives in tmp_path. No ONNX model, no sherpa-onnx import.
"""

import json
import logging
import os
import sys
import threading
import time
import types

import numpy as np
import pytest

from stt_daemon import events
from stt_daemon.config import Config
from stt_daemon.daemon import Daemon
from stt_daemon.embedder import SherpaEmbedder, build_embedder
from stt_daemon.segmenter import Utterance
from stt_daemon.transcriber import TranscriptResult

VEC_A = np.array([1.0, 0.0, 0.0], dtype=np.float32)
VEC_B = np.array([0.0, 1.0, 0.0], dtype=np.float32)


class FakeEmbedder:
    """Returns canned vectors; tracks peak concurrent embed() calls."""

    def __init__(self, vectors=None, delay=0.0, error=None):
        self._vectors = list(vectors) if vectors else []
        self._delay = delay
        self._error = error
        self.calls = 0
        self.max_in_flight = 0
        self._in_flight = 0
        self._lock = threading.Lock()

    def embed(self, audio):
        with self._lock:
            self.calls += 1
            self._in_flight += 1
            self.max_in_flight = max(self.max_in_flight, self._in_flight)
        try:
            if self._delay:
                time.sleep(self._delay)
            if self._error is not None:
                raise self._error
            if self._vectors:
                return np.array(self._vectors.pop(0), dtype=np.float32)
            return np.array(VEC_A, dtype=np.float32)
        finally:
            with self._lock:
                self._in_flight -= 1


class FakeTranscriber:
    def __init__(self, text="hello"):
        self.text = text

    def transcribe(self, audio):
        return TranscriptResult(language="en", language_probability=0.9, text=self.text)

    def close(self):
        pass


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

    def emit(self, result, duration_s, rms=0.0, reason="unknown", audio=None,
             transcribe_s=None, **kwargs):
        if not result.text:
            return None
        self.records.append(dict(kwargs))
        return None

    def close(self):
        pass


class FakeCapture:
    def frames(self):
        return iter(())

    def healthy(self):
        return True

    def close(self):
        pass


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


def _config(tmp_path, **overrides):
    cfg = Config()
    cfg.events.socket = str(tmp_path / "events.sock")
    cfg.output.stdout = False
    cfg.output.jsonl_dir = None
    cfg.output.save_audio = False
    cfg.speaker.enabled = True
    cfg.speaker.model_path = str(tmp_path / "model.onnx")
    cfg.speaker.store_dir = str(tmp_path / "speakers")
    for key, value in overrides.items():
        setattr(cfg.speaker, key, value)
    return cfg


def _daemon_with_lane(cfg, embedder, published=None):
    """Wire a daemon with fakes and a started speaker lane (no threads but lane)."""
    sink = FakeSink()
    daemon = Daemon(
        cfg,
        log_level=logging.WARNING,
        transcriber_factory=lambda: FakeTranscriber(),
        sink_factory=lambda: sink,
        embedder_factory=lambda: embedder,
    )
    daemon.transcriber = FakeTranscriber()
    daemon.sink = sink
    collected = [] if published is None else published
    daemon._publish = lambda type_, **fields: collected.append((type_, fields))
    daemon._start_speaker_lane()
    assert daemon._speaker_lane is not None
    daemon._speaker_lane.start()
    return daemon, sink, collected


def _transcripts(collected):
    return [fields for type_, fields in collected if type_ == events.TRANSCRIPT]


def _speakers_events(collected):
    return [fields for type_, fields in collected if type_ == events.SPEAKERS]


def test_no_sherpa_import_at_module_level():
    import stt_daemon.embedder  # noqa: F401 -- must import without the optional extra

    assert "sherpa_onnx" not in sys.modules


def test_build_embedder_none_when_disabled_or_missing(tmp_path, caplog):
    cfg = _config(tmp_path)
    cfg.speaker.enabled = False
    with caplog.at_level(logging.WARNING):
        assert build_embedder(cfg) is None
    assert "speaker" in caplog.text.lower()

    cfg.speaker.enabled = True
    cfg.speaker.model_path = None
    with caplog.at_level(logging.WARNING):
        assert build_embedder(cfg) is None

    cfg.speaker.model_path = str(tmp_path / "nope.onnx")
    with caplog.at_level(logging.WARNING):
        assert build_embedder(cfg) is None


def test_mint_then_matched_flow(tmp_path):
    cfg = _config(tmp_path)
    cfg.speaker.confirm_s = 100.0  # keep the third utterance at adapt, not confirm
    daemon, sink, collected = _daemon_with_lane(cfg, FakeEmbedder())

    daemon._process(_utt(1))
    daemon._process(_utt(2))
    daemon._process(_utt(3))
    daemon.stop()

    transcripts = _transcripts(collected)
    assert len(transcripts) == 3
    for fields in transcripts:
        for key in ("speaker_id", "speaker_status", "speaker", "speaker_score"):
            assert key in fields
        assert fields["seg_id"] in (1, 2, 3)

    assert transcripts[0]["speaker_status"] == "unknown"
    assert transcripts[0]["speaker_id"] is None

    assert transcripts[1]["speaker_status"] == "unknown"
    assert transcripts[1]["speaker_id"] == "spk_0001"
    assert transcripts[1]["speaker_score"] is None

    assert transcripts[2]["speaker_status"] == "matched"
    assert transcripts[2]["speaker_id"] == "spk_0001"
    assert transcripts[2]["speaker"] == "spk_0001"  # unnamed: id as label
    assert transcripts[2]["speaker_score"] == pytest.approx(1.0)

    # Sink saw the same five fields, including seg_id.
    assert sink.records[2]["seg_id"] == 3
    assert sink.records[2]["speaker_id"] == "spk_0001"
    assert sink.records[2]["speaker_status"] == "matched"

    # Mint publishes SPEAKERS; adapt does not.
    assert len(_speakers_events(collected)) == 1

    # Matched + named resolves the display name.
    assert daemon.speaker_store.rename("spk_0001", "Efren") is True
    daemon2, sink2, collected2 = _daemon_with_lane(cfg, FakeEmbedder())
    daemon2._process(_utt(4))
    daemon2.stop()
    assert _transcripts(collected2)[0]["speaker"] == "Efren"


def test_speakers_published_on_confirm_not_adapt(tmp_path):
    cfg = _config(tmp_path)
    cfg.speaker.confirm_s = 6.0
    daemon, _, collected = _daemon_with_lane(cfg, FakeEmbedder())

    daemon._process(_utt(1))
    daemon._process(_utt(2))  # mint (cumulative 3.0)
    assert len(_speakers_events(collected)) == 1
    daemon._process(_utt(3))  # adapt -> cumulative 6.0 -> confirm
    daemon.stop()

    assert len(_speakers_events(collected)) == 2
    assert _speakers_events(collected)[-1]["speakers"][0]["status"] == "confirmed"
    assert daemon.speaker_store.speakers["spk_0001"]["status"] == "confirmed"


def test_stalled_embedder_times_out_without_blocking(tmp_path):
    cfg = _config(tmp_path, join_timeout_s=0.2)
    daemon, sink, collected = _daemon_with_lane(cfg, FakeEmbedder(delay=2.0))

    started = time.monotonic()
    daemon._process(_utt(1))
    elapsed = time.monotonic() - started
    daemon.stop()

    assert elapsed < 2.0
    transcripts = _transcripts(collected)
    assert len(transcripts) == 1  # transcription still emitted
    assert transcripts[0]["speaker_status"] == "none"
    assert transcripts[0]["speaker_id"] is None
    assert sink.records[0]["speaker_status"] == "none"


def test_embedder_error_degrades_to_none(tmp_path):
    cfg = _config(tmp_path)
    daemon, _, collected = _daemon_with_lane(cfg, FakeEmbedder(error=RuntimeError("boom")))

    daemon._process(_utt(1))  # must not raise
    daemon.stop()

    transcripts = _transcripts(collected)
    assert len(transcripts) == 1
    assert transcripts[0]["speaker_status"] == "none"


@pytest.mark.parametrize("enabled,model_path", [(False, None), (True, None)])
def test_disabled_or_missing_model_has_no_lane(tmp_path, monkeypatch, enabled, model_path):
    monkeypatch.chdir(tmp_path)
    cfg = Config()
    cfg.output.stdout = False
    cfg.speaker.enabled = enabled
    cfg.speaker.model_path = model_path
    cfg.speaker.store_dir = "speakers"
    sink = FakeSink()
    daemon = Daemon(
        cfg,
        log_level=logging.WARNING,
        transcriber_factory=lambda: FakeTranscriber(),
        sink_factory=lambda: sink,
    )
    daemon.transcriber = FakeTranscriber()
    daemon.sink = sink
    collected = []
    daemon._publish = lambda type_, **fields: collected.append((type_, fields))
    daemon._start_speaker_lane()

    assert daemon._speaker_lane is None
    assert daemon.speaker_store is None

    daemon._process(_utt(1))
    transcripts = _transcripts(collected)
    assert transcripts[0]["speaker_status"] == "none"
    assert transcripts[0]["speaker_id"] is None
    assert transcripts[0]["speaker"] == "?"
    assert transcripts[0]["speaker_score"] is None
    assert not (tmp_path / "speakers").exists()


def test_inbox_valid_request_renames_and_publishes(tmp_path):
    cfg = _config(tmp_path)
    daemon, _, collected = _daemon_with_lane(cfg, FakeEmbedder())
    daemon._process(_utt(1))
    daemon._process(_utt(2))  # mint spk_0001
    before = len(_speakers_events(collected))

    inbox = tmp_path / "speakers" / "inbox"
    inbox.mkdir(parents=True, exist_ok=True)
    request = inbox / "req-1.json"
    request.write_text(json.dumps({"speaker_id": "spk_0001", "name": "Efren"}))

    daemon._speaker_lane.drain_inbox()

    assert daemon.speaker_store.speakers["spk_0001"]["name"] == "Efren"
    assert not request.exists()
    assert len(_speakers_events(collected)) == before + 1
    daemon.stop()


def test_inbox_malformed_and_unknown_id_quarantined(tmp_path, caplog):
    cfg = _config(tmp_path)
    daemon, _, collected = _daemon_with_lane(cfg, FakeEmbedder())
    before = len(_speakers_events(collected))

    inbox = tmp_path / "speakers" / "inbox"
    inbox.mkdir(parents=True, exist_ok=True)
    bad = inbox / "bad.json"
    bad.write_text("{not json")
    wrong_shape = inbox / "shape.json"
    wrong_shape.write_text(json.dumps({"speaker_id": "spk_0001"}))
    unknown = inbox / "unknown.json"
    unknown.write_text(json.dumps({"speaker_id": "spk_9999", "name": "Ghost"}))

    with caplog.at_level(logging.WARNING):
        daemon._speaker_lane.drain_inbox()
    daemon.stop()

    rejected = inbox / "rejected"
    assert (rejected / "bad.json").exists()
    assert (rejected / "shape.json").exists()
    assert (rejected / "unknown.json").exists()
    assert len(_speakers_events(collected)) == before  # no publish on reject
    assert "speaker" in caplog.text.lower()


def test_inbox_request_present_before_start_is_applied(tmp_path):
    cfg = _config(tmp_path)
    cfg.speaker.confirm_s = 100.0
    first, _, _ = _daemon_with_lane(cfg, FakeEmbedder())
    first._process(_utt(1))
    first._process(_utt(2))  # mint spk_0001, persisted
    first.stop()

    inbox = tmp_path / "speakers" / "inbox"
    inbox.mkdir(parents=True, exist_ok=True)
    (inbox / "pending.json").write_text(
        json.dumps({"speaker_id": "spk_0001", "name": "Efren"})
    )

    sink = FakeSink()
    daemon = Daemon(
        cfg,
        log_level=logging.WARNING,
        transcriber_factory=lambda: FakeTranscriber(),
        sink_factory=lambda: sink,
        embedder_factory=lambda: FakeEmbedder(),
    )
    daemon.transcriber = FakeTranscriber()
    daemon.sink = sink
    collected = []
    daemon._publish = lambda type_, **fields: collected.append((type_, fields))
    daemon._start_speaker_lane()
    daemon._speaker_lane.start()

    deadline = time.time() + 5.0
    while time.time() < deadline:
        if daemon.speaker_store.speakers.get("spk_0001", {}).get("name") == "Efren":
            break
        time.sleep(0.01)
    daemon.stop()

    assert daemon.speaker_store.speakers["spk_0001"]["name"] == "Efren"
    assert not (inbox / "pending.json").exists()


def test_embedding_lane_is_serial(tmp_path):
    cfg = _config(tmp_path)
    embedder = FakeEmbedder(delay=0.02)
    daemon, _, _ = _daemon_with_lane(cfg, embedder)
    lane = daemon._speaker_lane

    futures = [lane.submit(i, np.zeros(8, dtype=np.float32), 3.0, 0.05) for i in range(6)]
    for future in futures:
        future.result(timeout=10)
    daemon.stop()

    assert embedder.calls == 6
    assert embedder.max_in_flight == 1


# --- real SherpaEmbedder against a fake sherpa_onnx module --------------------


def _fake_sherpa_module(calls, vector, valid=True):
    """Fake sherpa_onnx implementing the verified API surface.

    Mirrors the official API: ``SpeakerEmbeddingExtractorConfig(model,
    num_threads=1, debug=False, provider="cpu")`` with ``validate()``, and
    ``SpeakerEmbeddingExtractor`` with ``create_stream / is_ready / compute``
    plus a readonly ``dim``. Unknown config kwargs raise TypeError, like the
    real pybind class.
    """
    module = types.ModuleType("sherpa_onnx")

    class Config:
        def __init__(self, model, num_threads=1, debug=False, provider="cpu"):
            calls.append(
                (
                    "config",
                    {
                        "model": model,
                        "num_threads": num_threads,
                        "debug": debug,
                        "provider": provider,
                    },
                )
            )

        def validate(self):
            calls.append(("validate",))
            return valid

    class Stream:
        def accept_waveform(self, sample_rate, waveform):
            calls.append(("accept_waveform", sample_rate, tuple(np.asarray(waveform).shape)))

        def input_finished(self):
            calls.append(("input_finished",))

    class Extractor:
        def __init__(self, config):
            calls.append(("extractor",))

        @property
        def dim(self):
            return len(vector)

        def create_stream(self):
            calls.append(("create_stream",))
            return Stream()

        def is_ready(self, stream):
            calls.append(("is_ready",))
            return True

        def compute(self, stream):
            calls.append(("compute",))
            return np.array(vector, dtype=np.float32)

    module.SpeakerEmbeddingExtractorConfig = Config
    module.SpeakerEmbeddingExtractor = Extractor
    return module


def test_sherpa_embedder_uses_verified_api(monkeypatch):
    calls = []
    vector = np.arange(192, dtype=np.float32)
    monkeypatch.setitem(sys.modules, "sherpa_onnx", _fake_sherpa_module(calls, vector))

    embedder = SherpaEmbedder("model.onnx")
    out = embedder.embed(np.zeros(1600, dtype=np.float32))

    assert out.shape == (192,)
    assert out.dtype == np.float32
    np.testing.assert_array_equal(out, vector)
    assert [call[0] for call in calls] == [
        "config",
        "validate",
        "extractor",
        "create_stream",
        "accept_waveform",
        "input_finished",
        "is_ready",
        "compute",
    ]
    assert calls[0][1] == {
        "model": "model.onnx",
        "num_threads": 1,
        "debug": False,
        "provider": "cpu",
    }
    assert calls[4][1] == 16000  # extractor runs at 16 kHz


def test_sherpa_embedder_invalid_config_raises(monkeypatch):
    calls = []
    monkeypatch.setitem(
        sys.modules, "sherpa_onnx", _fake_sherpa_module(calls, np.zeros(2), valid=False)
    )
    with pytest.raises(RuntimeError, match="failed to load"):
        SherpaEmbedder("model.onnx")


def test_sherpa_embedder_without_package_raises(monkeypatch):
    monkeypatch.setitem(sys.modules, "sherpa_onnx", None)  # import -> ImportError
    with pytest.raises(RuntimeError, match="sherpa-onnx is not installed"):
        SherpaEmbedder("model.onnx")


def test_sherpa_embedder_native_lib_failure_surfaces_cause(monkeypatch):
    monkeypatch.setitem(sys.modules, "sherpa_onnx", None)  # import -> ModuleNotFoundError
    with pytest.raises(RuntimeError, match="uv sync --extra speaker"):
        SherpaEmbedder("model.onnx")

    import builtins

    real_import = builtins.__import__

    def fake_import(name, *args, **kwargs):
        if name == "sherpa_onnx":
            raise ImportError(
                "libonnxruntime.so: cannot open shared object file: No such file or directory"
            )
        return real_import(name, *args, **kwargs)

    monkeypatch.delitem(sys.modules, "sherpa_onnx", raising=False)
    monkeypatch.setattr(builtins, "__import__", fake_import)
    with pytest.raises(RuntimeError, match="libonnxruntime.so"):
        SherpaEmbedder("model.onnx")


def _lane_without_thread(cfg, embedder, collected):
    """Build a daemon + lane without starting the worker (deterministic drain)."""
    sink = FakeSink()
    daemon = Daemon(
        cfg,
        log_level=logging.WARNING,
        transcriber_factory=lambda: FakeTranscriber(),
        sink_factory=lambda: sink,
        embedder_factory=lambda: embedder,
    )
    daemon.transcriber = FakeTranscriber()
    daemon.sink = sink
    daemon._publish = lambda type_, **fields: collected.append((type_, fields))
    daemon._start_speaker_lane()
    assert daemon._speaker_lane is not None
    return daemon


def test_inbox_fifo_quarantined_without_blocking(tmp_path):
    cfg = _config(tmp_path)
    collected: list = []
    daemon = _lane_without_thread(cfg, FakeEmbedder(), collected)
    lane = daemon._speaker_lane
    assert lane is not None
    try:
        inbox = tmp_path / "speakers" / "inbox"
        inbox.mkdir(parents=True, exist_ok=True)
        fifo = inbox / "evil.json"
        try:
            os.mkfifo(fifo)
        except (AttributeError, OSError, NotImplementedError) as e:
            pytest.skip(f"mkfifo unsupported: {e}")
        before = len(_speakers_events(collected))

        worker = threading.Thread(target=lane.drain_inbox)
        worker.start()
        worker.join(timeout=5)
        assert not worker.is_alive()  # must not block on the FIFO

        assert not os.path.lexists(fifo)
        assert os.path.lexists(inbox / "rejected" / "evil.json")
        assert len(_speakers_events(collected)) == before
    finally:
        daemon.stop()


def test_inbox_symlink_quarantined_target_untouched(tmp_path):
    cfg = _config(tmp_path)
    collected: list = []
    daemon = _lane_without_thread(cfg, FakeEmbedder(), collected)
    lane = daemon._speaker_lane
    assert lane is not None
    try:
        inbox = tmp_path / "speakers" / "inbox"
        inbox.mkdir(parents=True, exist_ok=True)
        target = tmp_path / "target.txt"
        target.write_text("secret", encoding="utf-8")
        link = inbox / "link.json"
        try:
            link.symlink_to(target)
        except (OSError, NotImplementedError) as e:
            pytest.skip(f"symlink unsupported: {e}")
        before = len(_speakers_events(collected))

        lane.drain_inbox()

        assert not os.path.lexists(link)
        assert os.path.lexists(inbox / "rejected" / "link.json")
        assert target.read_text(encoding="utf-8") == "secret"
        assert len(_speakers_events(collected)) == before
    finally:
        daemon.stop()


def test_inbox_oversized_request_quarantined(tmp_path):
    cfg = _config(tmp_path)
    collected: list = []
    daemon = _lane_without_thread(cfg, FakeEmbedder(), collected)
    lane = daemon._speaker_lane
    assert lane is not None
    try:
        inbox = tmp_path / "speakers" / "inbox"
        inbox.mkdir(parents=True, exist_ok=True)
        big = inbox / "big.json"
        big.write_bytes(b"x" * (65536 + 1024))
        before = len(_speakers_events(collected))

        lane.drain_inbox()

        assert not big.exists()
        assert (inbox / "rejected" / "big.json").exists()
        assert len(_speakers_events(collected)) == before
    finally:
        daemon.stop()