STT-Assistant/tests/test_evaluate.py

"""Hardware/network-free tests for stt_daemon.evaluate."""

import wave

import numpy as np

from stt_daemon.evaluate import EvalResult, accuracy, evaluate, format_report
from stt_daemon.speakers import SpeakerStore

_SAMPLE_RATE = 16000

ALICE = np.array([1.0, 0.0, 0.0], dtype=np.float32)
BOB = np.array([0.0, 1.0, 0.0], dtype=np.float32)
# Cosine 0.5 to ALICE, 0.0 to BOB -> uncertain with top-1 alice (default thresholds).
ALICE_UNCERTAIN = np.array([0.5, 0.0, 0.8660254], dtype=np.float32)
UNKNOWN_VEC = np.array([0.0, 0.0, 1.0], dtype=np.float32)

MARKERS = {
    1001: ALICE,
    1002: BOB,
    1003: ALICE_UNCERTAIN,
    1004: UNKNOWN_VEC,
}


class FakeEmbedder:
    """Return a controlled unit vector keyed by the first int16 sample."""

    def __init__(self, mapping=None):
        self.mapping = dict(MARKERS if mapping is None else mapping)

    def embed(self, audio):
        marker = int(round(float(np.asarray(audio).reshape(-1)[0]) * 32768.0))
        return np.array(self.mapping[marker], dtype=np.float32)


def _write_wav(path, marker, nframes=1600):
    path.parent.mkdir(parents=True, exist_ok=True)
    frames = np.zeros(nframes, dtype=np.int16)
    frames[0] = np.int16(marker)
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(_SAMPLE_RATE)
        wav.writeframes(frames.tobytes())


def _write_wav_rate(path, marker, framerate):
    path.parent.mkdir(parents=True, exist_ok=True)
    frames = np.full(1600, marker, dtype=np.int16)
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(framerate)
        wav.writeframes(frames.tobytes())


def _write_stereo_wav(path):
    path.parent.mkdir(parents=True, exist_ok=True)
    frames = np.zeros((1600, 2), dtype=np.int16)
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(2)
        wav.setsampwidth(2)
        wav.setframerate(_SAMPLE_RATE)
        wav.writeframes(frames.tobytes())


def _make_store(root):
    store = SpeakerStore(root)
    store.enroll("alice", [ALICE])
    store.enroll("bob", [BOB])
    return store


def test_all_correct_accuracy_one(tmp_path):
    store = _make_store(tmp_path / "store")
    labeled = tmp_path / "labeled"
    _write_wav(labeled / "alice" / "a.wav", 1001)
    _write_wav(labeled / "bob" / "b.wav", 1002)
    result = evaluate(store, FakeEmbedder(), labeled)
    assert result.total == 2
    assert result.correct == 2
    assert result.wrong == 0
    assert result.unknown == 0
    assert accuracy(result) == 1.0
    assert result.per_speaker == {"alice": (1, 1), "bob": (1, 1)}
    assert result.confusions == {}
    assert len(result.margins) == 2
    assert result.labeled == (("matched", "alice/a.wav"), ("matched", "bob/b.wav"))


def test_wrong_match_populates_confusions(tmp_path):
    store = _make_store(tmp_path / "store")
    labeled = tmp_path / "labeled"
    _write_wav(labeled / "alice" / "a.wav", 1002)  # bob's voice under alice/
    result = evaluate(store, FakeEmbedder(), labeled)
    assert result.total == 1
    assert result.correct == 0
    assert result.wrong == 1
    assert result.confusions == {("alice", "bob"): 1}
    assert result.per_speaker == {"alice": (1, 0)}
    assert accuracy(result) == 0.0


def test_uncertain_top1_correct(tmp_path):
    store = _make_store(tmp_path / "store")
    labeled = tmp_path / "labeled"
    _write_wav(labeled / "alice" / "u.wav", 1003)
    result = evaluate(store, FakeEmbedder(), labeled)
    assert result.total == 1
    assert result.uncertain_correct == 1
    assert result.uncertain_wrong == 0
    assert result.correct == 0 and result.wrong == 0 and result.unknown == 0
    assert result.labeled[0][0] == "uncertain"
    assert len(result.margins) == 1


def test_uncertain_top1_wrong_and_unknown(tmp_path):
    store = _make_store(tmp_path / "store")
    labeled = tmp_path / "labeled"
    # Uncertain vector closest to alice but filed under bob -> uncertain_wrong.
    _write_wav(labeled / "bob" / "u.wav", 1003)
    _write_wav(labeled / "alice" / "x.wav", 1004)
    result = evaluate(store, FakeEmbedder(), labeled)
    assert result.uncertain_wrong == 1
    assert result.unknown == 1
    assert result.total == 2


def test_unreadable_and_bad_format_skipped(tmp_path):
    store = _make_store(tmp_path / "store")
    labeled = tmp_path / "labeled"
    _write_wav(labeled / "alice" / "good.wav", 1001)
    (labeled / "alice" / "bad.wav").parent.mkdir(parents=True, exist_ok=True)
    (labeled / "alice" / "bad.wav").write_bytes(b"not a wav at all")
    _write_wav_rate(labeled / "alice" / "rate8k.wav", 1001, 8000)
    _write_stereo_wav(labeled / "alice" / "stereo.wav")
    result = evaluate(store, FakeEmbedder(), labeled)
    assert result.total == 1
    assert result.correct == 1
    assert result.labeled == (("matched", "alice/good.wav"),)


def test_unknown_speaker_directory_still_evaluated(tmp_path):
    store = _make_store(tmp_path / "store")
    labeled = tmp_path / "labeled"
    _write_wav(labeled / "carol" / "c.wav", 1001)  # matches alice, true name carol
    result = evaluate(store, FakeEmbedder(), labeled)
    assert result.total == 1
    assert result.wrong == 1
    assert result.confusions == {("carol", "alice"): 1}
    assert result.per_speaker == {"carol": (1, 0)}


def test_empty_and_missing_directory(tmp_path):
    store = _make_store(tmp_path / "store")
    empty = tmp_path / "empty"
    empty.mkdir()
    result = evaluate(store, FakeEmbedder(), empty)
    assert result.total == 0
    assert accuracy(result) == 0.0
    assert result.margins == ()
    assert result.labeled == ()
    missing = tmp_path / "does-not-exist"
    result2 = evaluate(store, FakeEmbedder(), missing)
    assert result2.total == 0
    assert accuracy(result2) == 0.0


def test_format_report_contents_and_order(tmp_path):
    result = EvalResult(
        total=3,
        correct=1,
        wrong=1,
        uncertain_correct=1,
        uncertain_wrong=0,
        unknown=0,
        per_speaker={"bob": (2, 1), "alice": (1, 0)},
        confusions={("alice", "bob"): 1},
        margins=(0.5, 1.0),
        labeled=(("matched", "alice/a.wav"),),
    )
    report = format_report(result)
    assert "total: 3" in report
    assert "accuracy: 0.333" in report
    assert "correct: 1" in report
    assert "wrong: 1" in report
    assert "uncertain (top-1 correct/total): 1/1" in report
    assert "unknown: 0" in report
    assert "alice | 1 | 0 | 0.000" in report
    assert "bob | 2 | 1 | 0.500" in report
    assert report.index("alice |") < report.index("bob |")
    assert "alice -> bob  (1)" in report
    assert "count: 2" in report
    assert "mean: 0.7500" in report
    assert "p10:" in report
    # Empty margins guard.
    empty_report = format_report(
        EvalResult(0, 0, 0, 0, 0, 0, {}, {}, (), ()),
    )
    assert "count: 0" in empty_report


def test_margins_nonempty_and_store_untouched(tmp_path):
    store = _make_store(tmp_path / "store")
    before_speakers = store.speakers
    before_vectors = np.array(store._vectors, copy=True)
    labeled = tmp_path / "labeled"
    _write_wav(labeled / "alice" / "a.wav", 1001)
    result = evaluate(store, FakeEmbedder(), labeled)
    assert len(result.margins) == 1
    assert result.margins[0] > 0.0
    assert store.speakers == before_speakers
    assert np.array_equal(np.asarray(store._vectors), before_vectors)