STT-Assistant/tests/test_speaker_enroll.py

"""Tests for offline speaker pre-enrollment (WAV reading + CLI handler).

Hardware- and network-free: WAVs are synthesized with stdlib ``wave`` and the
embedder is a fake; the real sherpa-onnx package is never required.
"""

import argparse
import wave
from pathlib import Path

import numpy as np
import pytest

import stt_daemon.cli as cli_module
from stt_daemon.enroll import enroll_from_files, read_wav
from stt_daemon.speakers import SpeakerStore


def write_mono_wav(path: Path, samples: np.ndarray, sample_rate: int = 16000,
                   nchannels: int = 1) -> None:
    """Write int16 PCM WAV frames from float32 samples in [-1, 1]."""
    frames = (np.asarray(samples, dtype=np.float32) * 32767.0).astype(np.int16)
    if nchannels == 2:
        frames = np.stack([frames, frames], axis=1).reshape(-1)
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(nchannels)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(frames.tobytes())


def make_wav(path: Path, seconds: float, sample_rate: int = 16000,
             nchannels: int = 1) -> Path:
    """Write a sine clip of ``seconds`` and return the path."""
    count = int(seconds * sample_rate)
    t = np.arange(count, dtype=np.float32) / sample_rate
    samples = 0.5 * np.sin(2.0 * np.pi * 440.0 * t).astype(np.float32)
    write_mono_wav(path, samples, sample_rate, nchannels)
    return path


class FakeEmbedder:
    """Deterministic unit-norm embedding from audio energy."""

    def __init__(self) -> None:
        self.calls = 0

    def embed(self, audio: np.ndarray) -> np.ndarray:
        self.calls += 1
        energy = float(np.mean(np.asarray(audio, dtype=np.float32) ** 2))
        vector = np.array([energy + self.calls, 1.0], dtype=np.float32)
        return vector / np.linalg.norm(vector)


# -- read_wav -----------------------------------------------------------------


def test_read_wav_happy_path(tmp_path):
    wav = make_wav(tmp_path / "clip.wav", seconds=1.0)

    samples, rate = read_wav(wav)

    assert rate == 16000
    assert samples.dtype == np.float32
    assert len(samples) == 16000
    assert float(np.abs(samples).max()) <= 1.0


def test_read_wav_rejects_stereo(tmp_path):
    wav = make_wav(tmp_path / "stereo.wav", seconds=1.0, nchannels=2)

    with pytest.raises(ValueError, match="mono"):
        read_wav(wav)


def test_read_wav_rejects_non_16k(tmp_path):
    wav = make_wav(tmp_path / "clip8k.wav", seconds=1.0, sample_rate=8000)

    with pytest.raises(ValueError, match="16 kHz"):
        read_wav(wav)


def test_read_wav_rejects_non_wav(tmp_path):
    bogus = tmp_path / "notes.txt"
    bogus.write_text("not audio", encoding="utf-8")

    with pytest.raises(ValueError, match="not a WAV"):
        read_wav(bogus)


def test_read_wav_rejects_unsupported_width(tmp_path):
    wav = tmp_path / "clip8bit.wav"
    with wave.open(str(wav), "wb") as handle:
        handle.setnchannels(1)
        handle.setsampwidth(1)
        handle.setframerate(16000)
        handle.writeframes(bytes([128] * 16000))

    with pytest.raises(ValueError, match="sample width"):
        read_wav(wav)


# -- enroll_from_files ----------------------------------------------------------


def test_enroll_from_files_enrolls_and_saves(tmp_path):
    wav = make_wav(tmp_path / "clip.wav", seconds=2.5)
    store = SpeakerStore(tmp_path / "store")
    embedder = FakeEmbedder()

    speaker_id = enroll_from_files(store, embedder, "Efren", [wav], 2.0)

    assert speaker_id == "spk_0001"
    assert embedder.calls == 1
    assert (tmp_path / "store" / "speakers.json").exists()
    fresh = SpeakerStore(tmp_path / "store")
    fresh.load()
    assert fresh.speakers[speaker_id]["name"] == "Efren"
    assert fresh.speakers[speaker_id]["status"] == "confirmed"


def test_enroll_from_files_rejects_short_clip(tmp_path):
    wav = make_wav(tmp_path / "short.wav", seconds=0.5)
    store = SpeakerStore(tmp_path / "store")

    with pytest.raises(ValueError, match="too short"):
        enroll_from_files(store, FakeEmbedder(), "Efren", [wav], 2.0)


# -- CLI handler ------------------------------------------------------------------


def run_handler(monkeypatch, tmp_path, *, name="Efren", wavs=None,
                model_ok=True, min_enroll_s=2.0, force=False):
    """Drive ``_handle_speaker_enroll`` with fakes; return (rc, out, err)."""
    if wavs is None:
        wavs = [make_wav(tmp_path / "clip.wav", seconds=2.5)]
    model = tmp_path / "model.onnx"
    if model_ok:
        model.write_bytes(b"fake-model")
    cfg_path = tmp_path / "config.toml"
    cfg_path.write_text(
        "[speaker]\n"
        f'model_path = "{model}"\n'
        f"min_enroll_s = {min_enroll_s}\n"
        f'store_dir = "{tmp_path / "speakers"}"\n'
        "[events]\n"
        f'socket = "{tmp_path / "bus.sock"}"\n',
        encoding="utf-8",
    )

    fake = FakeEmbedder()

    class FakeSherpa:
        def __init__(self, model_path, provider="cpu"):
            self.embedder = fake

        def embed(self, audio):
            return fake.embed(audio)

    monkeypatch.setattr(cli_module, "SherpaEmbedder", FakeSherpa)
    args = argparse.Namespace(
        config=str(cfg_path), name=name, wav=[str(w) for w in wavs], force=force
    )
    import io
    from contextlib import redirect_stderr, redirect_stdout

    out, err = io.StringIO(), io.StringIO()
    with redirect_stdout(out), redirect_stderr(err):
        rc = cli_module._handle_speaker_enroll(args)
    return rc, out.getvalue(), err.getvalue()


def test_cli_handler_success(monkeypatch, tmp_path):
    rc, out, _ = run_handler(monkeypatch, tmp_path)

    assert rc == 0
    assert "enrolled 'Efren' as spk_0001" in out
    assert "1 clip(s)" in out


def test_cli_handler_missing_model(monkeypatch, tmp_path):
    rc, _, err = run_handler(monkeypatch, tmp_path, model_ok=False)

    assert rc == 1
    assert "model" in err.lower()


def test_cli_handler_short_wav(monkeypatch, tmp_path):
    short = make_wav(tmp_path / "short.wav", seconds=0.5)

    rc, _, err = run_handler(monkeypatch, tmp_path, wavs=[short])

    assert rc == 1
    assert "too short" in err


def test_cli_handler_dedupes_and_saves(monkeypatch, tmp_path):
    """Two identical clips collapse to one prototype; the line says so."""
    wavs = [make_wav(tmp_path / f"clip{i}.wav", seconds=2.5) for i in range(2)]
    model = tmp_path / "model.onnx"
    model.write_bytes(b"fake-model")
    cfg_path = tmp_path / "config.toml"
    cfg_path.write_text(
        "[speaker]\n"
        f'model_path = "{model}"\n'
        "min_enroll_s = 2.0\n"
        f'store_dir = "{tmp_path / "speakers"}"\n'
        "[events]\n"
        f'socket = "{tmp_path / "bus.sock"}"\n',
        encoding="utf-8",
    )

    class ConstEmbedder:
        def embed(self, audio):
            return np.array([1.0, 0.0], dtype=np.float32)

    class FakeSherpa:
        def __init__(self, model_path, provider="cpu"):
            self._embedder = ConstEmbedder()

        def embed(self, audio):
            return self._embedder.embed(audio)

    monkeypatch.setattr(cli_module, "SherpaEmbedder", FakeSherpa)
    args = argparse.Namespace(
        config=str(cfg_path), name="Efren",
        wav=[str(w) for w in wavs], force=True, dedupe_threshold=0.98,
    )
    import io
    from contextlib import redirect_stderr, redirect_stdout

    out, err = io.StringIO(), io.StringIO()
    with redirect_stdout(out), redirect_stderr(err):
        rc = cli_module._handle_speaker_enroll(args)

    assert rc == 0
    assert "removed 1 duplicate prototype(s)" in out.getvalue()
    fresh = SpeakerStore(tmp_path / "speakers")
    fresh.load()
    assert fresh.speakers["spk_0001"]["prototypes"] == 1