STT-Assistant/tests/test_transcriber.py

"""Tests for the Whisper segment confidence gate."""

import numpy as np
import pytest

from stt_daemon import transcriber as transcriber_mod
from stt_daemon.transcriber import Transcriber, accept_segment


class _Seg:
    def __init__(self, text="hello", no_speech_prob=0.1, avg_logprob=-0.3):
        self.text = text
        self.no_speech_prob = no_speech_prob
        self.avg_logprob = avg_logprob


def test_accepts_confident_speech():
    assert accept_segment(_Seg(), 0.6, -1.0)


def test_rejects_high_no_speech_prob():
    assert not accept_segment(_Seg(no_speech_prob=0.9), 0.6, -1.0)


def test_rejects_low_avg_logprob():
    assert not accept_segment(_Seg(avg_logprob=-1.5), 0.6, -1.0)


def test_threshold_boundary_is_accepted():
    assert accept_segment(_Seg(no_speech_prob=0.6, avg_logprob=-1.0), 0.6, -1.0)


def test_missing_attributes_default_to_accept():
    class Bare:
        text = "x"

    assert accept_segment(Bare(), 0.6, -1.0)


class _Info:
    language = "en"
    language_probability = 1.0


class _FakeWhisperModel:
    """Records constructor/transcribe kwargs without loading a real model."""

    last = None

    def __init__(self, model, device, compute_type, cpu_threads=0):
        self.init_kwargs = {
            "model": model,
            "device": device,
            "compute_type": compute_type,
            "cpu_threads": cpu_threads,
        }
        self.transcribe_kwargs = None
        _FakeWhisperModel.last = self

    def transcribe(self, audio, **kwargs):
        self.transcribe_kwargs = kwargs
        return [], _Info()


@pytest.fixture
def fake_model(monkeypatch):
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _FakeWhisperModel)
    return _FakeWhisperModel


def _audio():
    return np.zeros(16000, dtype=np.float32)


def test_cpu_threads_forwarded_to_model(fake_model):
    t = Transcriber(
        "small", device="cpu", compute_type="int8", language=None, cpu_threads=16
    )
    assert t.cpu_threads == 16
    assert fake_model.last.init_kwargs["cpu_threads"] == 16


def test_cpu_threads_default_is_auto(fake_model):
    Transcriber("small", device="cpu", compute_type="int8", language=None)
    assert fake_model.last.init_kwargs["cpu_threads"] == 0


def test_transcribe_disables_timestamps(fake_model):
    t = Transcriber("small", device="cpu", compute_type="int8", language=None)
    t.transcribe(_audio())
    kwargs = fake_model.last.transcribe_kwargs
    assert kwargs["without_timestamps"] is True
    assert kwargs["condition_on_previous_text"] is False


def test_auto_language_stays_none_for_bilingual(fake_model):
    """language=auto must reach the engine as None so es/en both keep working."""
    t = Transcriber("small", device="cpu", compute_type="int8", language=None)
    t.transcribe(_audio())
    assert fake_model.last.transcribe_kwargs["language"] is None


def test_explicit_language_is_preserved(fake_model):
    t = Transcriber("small", device="cpu", compute_type="int8", language="es")
    t.transcribe(_audio())
    assert fake_model.last.transcribe_kwargs["language"] == "es"