STT-Assistant/tests/test_cli_speaker_eval.py

"""Tests for the speaker-eval CLI plus the new enroll threshold flags.

Hardware- and network-free: WAVs are synthesized with stdlib ``wave`` (a DC
level encodes the "speaker id"), the embedder is a fake, and the real
sherpa-onnx package is never required.
"""

import argparse
import io
import wave
from contextlib import redirect_stderr, redirect_stdout
from pathlib import Path

import numpy as np

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

A = 0.30
B = -0.30


def write_const_wav(path: Path, seconds: float, value: float) -> Path:
    """Write a constant-level int16 PCM WAV; the level encodes the speaker."""
    count = int(seconds * 16000)
    samples = np.full(count, value, dtype=np.float32)
    frames = (samples * 32767.0).astype(np.int16)
    with wave.open(str(path), "wb") as handle:
        handle.setnchannels(1)
        handle.setsampwidth(2)
        handle.setframerate(16000)
        handle.writeframes(frames.tobytes())
    return path


class LevelEmbedder:
    """Deterministic embedding from the clip DC level (the fake speaker id)."""

    def __init__(self, dim: int = 2) -> None:
        self.dim = dim

    def embed(self, audio: np.ndarray) -> np.ndarray:
        mean = float(np.mean(np.asarray(audio, dtype=np.float32)))
        if self.dim == 2:
            vector = [1.0, 0.0] if mean > 0.15 else (
                [0.0, 1.0] if mean < -0.15 else [0.7071, 0.7071])
        else:
            vector = [1.0, 0.0, 0.0] if mean > 0.15 else [0.0, 1.0, 0.0]
        return np.array(vector, dtype=np.float32)


class RotEmbedder:
    """Same-speaker clips with distinct-but-similar embeddings (20 deg/call)."""

    def embed(self, audio: np.ndarray) -> np.ndarray:
        mean = float(np.mean(np.asarray(audio, dtype=np.float32)))
        if mean < -0.15:
            return np.array([-1.0, 0.0], dtype=np.float32)
        RotEmbedder.calls = getattr(RotEmbedder, "calls", 0) + 1
        angle = np.deg2rad(20.0 * RotEmbedder.calls)
        vector = np.array([np.cos(angle), np.sin(angle)], dtype=np.float32)
        return vector / np.linalg.norm(vector)


def write_cfg(tmp_path: Path, *, model_ok: bool = True) -> Path:
    """Write a temp config with a temp speaker store; return the cfg path."""
    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'store_dir = "{tmp_path / "speakers"}"\n'
        "[events]\n"
        f'socket = "{tmp_path / "bus.sock"}"\n',
        encoding="utf-8",
    )
    return cfg_path


def patch_embedder(monkeypatch, embedder) -> None:
    """Point cli.SherpaEmbedder at a fake wrapping ``embedder``."""

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

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

    monkeypatch.setattr(cli_module, "SherpaEmbedder", FakeSherpa)


def enroll_name(store_dir: Path, name: str, clip: Path,
                embedder, min_enroll_s: float = 2.0) -> None:
    """Enroll one clip under ``name`` into the temp store."""
    store = SpeakerStore(store_dir)
    enroll_from_files(store, embedder, name, [clip], min_enroll_s)


def run_eval(monkeypatch, tmp_path, *, labeled, min_accuracy=None,
             model_ok=True, dim=2, pre_enroll_dim=None):
    """Drive ``_handle_speaker_eval`` with fakes; return (rc, out, err)."""
    cfg_path = write_cfg(tmp_path, model_ok=model_ok)
    if pre_enroll_dim is not None:
        clip = write_const_wav(tmp_path / "enroll.wav", 2.5, A)
        enroll_name(tmp_path / "speakers", "Efren", clip,
                    LevelEmbedder(dim=pre_enroll_dim))
    patch_embedder(monkeypatch, LevelEmbedder(dim=dim))
    args = argparse.Namespace(
        config=str(cfg_path), labeled=str(labeled), min_accuracy=min_accuracy)
    out, err = io.StringIO(), io.StringIO()
    with redirect_stdout(out), redirect_stderr(err):
        rc = cli_module._handle_speaker_eval(args)
    return rc, out.getvalue(), err.getvalue()


def make_labeled(tmp_path: Path, *, value: float = A,
                 name: str = "Efren") -> Path:
    """Build labeled/<name>/clip.wav at the given DC level."""
    labeled = tmp_path / "labeled"
    clip_dir = labeled / name
    clip_dir.mkdir(parents=True, exist_ok=True)
    write_const_wav(clip_dir / "clip.wav", 2.5, value)
    return labeled


def test_eval_report_success(monkeypatch, tmp_path):
    clip = write_const_wav(tmp_path / "enroll.wav", 2.5, A)
    enroll_name(tmp_path / "speakers", "Efren", clip, LevelEmbedder(dim=2))
    labeled = make_labeled(tmp_path)

    rc, out, _ = run_eval(monkeypatch, tmp_path, labeled=labeled)

    assert rc == 0
    assert "total: 1" in out
    assert "accuracy: 1.000" in out


def test_eval_below_min_accuracy(monkeypatch, tmp_path):
    clip = write_const_wav(tmp_path / "enroll.wav", 2.5, A)
    enroll_name(tmp_path / "speakers", "Efren", clip, LevelEmbedder(dim=2))
    labeled = make_labeled(tmp_path, value=B)

    rc, _, err = run_eval(monkeypatch, tmp_path, labeled=labeled,
                          min_accuracy=0.5)

    assert rc == 1
    assert "below --min-accuracy" in err


def test_eval_missing_model(monkeypatch, tmp_path):
    labeled = make_labeled(tmp_path)

    rc, _, err = run_eval(monkeypatch, tmp_path, labeled=labeled,
                          model_ok=False)

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


def test_eval_dim_mismatch(monkeypatch, tmp_path):
    labeled = make_labeled(tmp_path)

    rc, _, err = run_eval(monkeypatch, tmp_path, labeled=labeled,
                          dim=2, pre_enroll_dim=3)

    assert rc == 1
    assert "different model" in err


def test_eval_missing_labeled_path(monkeypatch, tmp_path):
    write_cfg(tmp_path)
    patch_embedder(monkeypatch, LevelEmbedder(dim=2))
    args = argparse.Namespace(
        config=str(tmp_path / "config.toml"),
        labeled=str(tmp_path / "does-not-exist"), min_accuracy=None)
    out, err = io.StringIO(), io.StringIO()
    with redirect_stdout(out), redirect_stderr(err):
        rc = cli_module._handle_speaker_eval(args)

    assert rc == 0
    assert "total: 0" in out.getvalue()


def test_seed_handler_forwards_new_thresholds(monkeypatch, tmp_path):
    audio_dir = tmp_path / "audio"
    audio_dir.mkdir(exist_ok=True)
    seed = write_const_wav(audio_dir / "seed.wav", 2.5, A)
    write_const_wav(audio_dir / "a1.wav", 2.5, A)
    write_const_wav(audio_dir / "b1.wav", 2.5, B)
    cfg_path = write_cfg(tmp_path)
    RotEmbedder.calls = 0
    patch_embedder(monkeypatch, RotEmbedder())
    args = argparse.Namespace(
        config=str(cfg_path), name="Efren", seed=str(seed), count=1,
        min_score=0.55, audio_dir=str(audio_dir), max_scan=0, yes=True,
        no_play=True, force=True, dedupe_threshold=0.97,
        other_threshold=0.44)
    out, err = io.StringIO(), io.StringIO()
    with redirect_stdout(out), redirect_stderr(err):
        rc = cli_module._handle_speaker_enroll_seed(args)

    assert rc == 0
    assert "from seed" in out.getvalue()