ApiCore/ApiCore-Whisper/tests/test_transcribe.py

"""Tests for transcribe.py with a fake model (no model files, no network)."""

from __future__ import annotations

import json
import os
import socket
import stat
import sys
import wave
from pathlib import Path
from types import SimpleNamespace

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import transcribe  # noqa: E402

SAMPLE_RATE = 8000


FORMATS = {"webm": ("webm", "libopus"), "ogg": ("ogg", "libopus"), "mp4": ("mp4", "aac")}
RATE = 48000
FRAME = 960


def write_wav(path: Path, seconds: float) -> None:
    """Writes silence as real audio in the container the extension names."""
    import av
    import numpy as np

    container_name, codec = FORMATS[path.suffix.lstrip(".")]
    out = av.open(str(path), "w", format=container_name)
    stream = out.add_stream(codec, rate=RATE)
    stream.layout = "mono"
    frame = av.AudioFrame.from_ndarray(np.zeros((1, FRAME), dtype=np.float32), format="flt", layout="mono")
    frame.sample_rate = RATE
    for index in range(int(seconds * RATE / FRAME)):
        frame.pts = index * FRAME
        for packet in stream.encode(frame):
            out.mux(packet)
    for packet in stream.encode(None):
        out.mux(packet)
    out.close()


class FakeModel:
    """Stands in for WhisperModel. ``texts`` is consumed one entry per call, so
    the n-th transcribed file gets the n-th entry (the audio is an ndarray)."""

    def __init__(self, texts=None, language="de", probability=1.0):
        self.calls: list[int] = []
        self.kwargs: list[dict] = []
        self.audio: list = []
        self.texts = list(texts or [])
        self.language = language
        self.probability = probability

    def transcribe(self, audio, **kwargs):
        self.calls.append(len(self.audio))
        self.audio.append(audio)
        self.kwargs.append(kwargs)
        entry = self.texts[len(self.calls) - 1] if len(self.calls) <= len(self.texts) else []
        segments = [
            SimpleNamespace(start=index * 2.0, end=index * 2.0 + 1.5, text=f" {text} ")
            for index, text in enumerate(entry)
        ]
        info = SimpleNamespace(language=self.language, language_probability=self.probability)
        return iter(segments), info


@pytest.fixture
def work(tmp_path: Path) -> Path:
    (tmp_path / "in").mkdir(mode=0o700)
    (tmp_path / "out").mkdir(mode=0o700)
    return tmp_path


def make_manifest(work: Path, names: list[str]) -> list[dict]:
    items = []
    for name in names:
        capture_id = name.split(".")[0]
        items.append({"captureId": capture_id, "file": name})
    manifest = {
        "v": 1,
        "items": [
            {"captureId": i["captureId"], "claimToken": "tok", "file": i["file"],
             "mime": "audio/webm", "reportedSec": None}
            for i in items
        ],
    }
    (work / "in" / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
    return items


def read_result(work: Path, capture_id: str) -> dict:
    return json.loads((work / "out" / f"{capture_id}.json").read_text(encoding="utf-8"))


def run_all(work: Path, model: FakeModel, names: list[str]) -> dict:
    items = make_manifest(work, names)
    return transcribe.run(work, lambda: model, "Durchsicht: Koenigin.", items)


def test_writes_an_empty_result_for_silence(work: Path):
    write_wav(work / "in" / "cap1.webm", 3)
    counters = run_all(work, FakeModel(), ["cap1.webm"])
    result = read_result(work, "cap1")
    assert result["error"] == "NO_SPEECH"
    assert result["text"] == ""
    assert result["segments"] == []
    assert counters["no_speech"] == 1


def test_rejects_audio_longer_than_240_seconds(work: Path):
    write_wav(work / "in" / "long1.webm", 241)
    model = FakeModel(texts=[["das darf nie ankommen"]])
    counters = run_all(work, model, ["long1.webm"])
    result = read_result(work, "long1")
    assert result["error"] == "TOO_LONG"
    assert result["durationSec"] > 240
    assert model.calls == []  # never reaches the model
    assert counters["too_long"] == 1


def test_accepts_audio_of_exactly_the_limit_region(work: Path):
    write_wav(work / "in" / "ok1.webm", 239)
    model = FakeModel(texts=[["Die Koenigin ist da"]])
    run_all(work, model, ["ok1.webm"])
    assert "error" not in read_result(work, "ok1")


def test_rejects_a_corrupt_file_without_crashing_the_batch(work: Path):
    (work / "in" / "bad1.webm").write_bytes(os.urandom(4096))
    write_wav(work / "in" / "good1.webm", 5)
    model = FakeModel(texts=[["Zwei Brutwaben mit Stiften"]])
    counters = run_all(work, model, ["bad1.webm", "good1.webm"])
    assert read_result(work, "bad1")["error"] == "DECODE_FAILED"
    assert read_result(work, "good1")["text"] == "Zwei Brutwaben mit Stiften"
    assert len(model.calls) == 1
    assert counters["failed"] == 1 and counters["transcribed"] == 1


def test_a_model_error_on_one_file_does_not_stop_the_others(work: Path):
    class Exploding(FakeModel):
        def transcribe(self, audio, **kwargs):
            if not self.calls and not self.audio:
                self.audio.append(None)
                self.calls.append(0)
                raise RuntimeError("secret transcript text must not leak")
            return super().transcribe(audio, **kwargs)

    write_wav(work / "in" / "a1.webm", 5)
    write_wav(work / "in" / "b1.webm", 5)
    model = Exploding(texts=[[], ["Volk ist ruhig und sanftmuetig"]])
    run_all(work, model, ["a1.webm", "b1.webm"])
    assert read_result(work, "a1")["error"] == "DECODE_FAILED"
    assert read_result(work, "b1")["text"].startswith("Volk")


def test_processes_files_one_after_the_other_and_writes_one_result_each(work: Path):
    names = ["n1.webm", "n2.webm", "n3.webm"]
    for name in names:
        write_wav(work / "in" / name, 4)
    model = FakeModel(texts=[[f"Notiz fuer {name}"] for name in names])
    counters = run_all(work, model, names)
    assert len(model.calls) == 3
    for name in names:
        assert read_result(work, name.split(".")[0])["text"] == f"Notiz fuer {name}"
    assert sorted(p.name for p in (work / "out").iterdir()) == ["n1.json", "n2.json", "n3.json"]
    assert counters["transcribed"] == 3


def test_loads_the_model_once_for_the_whole_batch(work: Path):
    names = ["m1.webm", "m2.webm"]
    for name in names:
        write_wav(work / "in" / name, 4)
    loads = []
    model = FakeModel(texts=[["Genug Text hier"] for _ in names])

    def factory():
        loads.append(1)
        return model

    transcribe.run(work, factory, "p", make_manifest(work, names))
    assert len(loads) == 1


def test_does_not_read_the_network(monkeypatch):
    monkeypatch.setattr(socket, "socket", socket.socket)  # restored after the test
    transcribe.block_network()
    with pytest.raises(OSError):
        socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    with pytest.raises(OSError):
        socket.create_connection(("127.0.0.1", 9), timeout=1)


def test_the_network_block_is_a_subclassable_class(monkeypatch):
    monkeypatch.setattr(socket, "socket", socket.socket)
    transcribe.block_network()

    class Later(socket.socket):  # what ssl does at import time
        pass

    with pytest.raises(OSError):
        Later()


def test_writes_segments_with_start_end_and_text(work: Path):
    write_wav(work / "in" / "seg1.webm", 10)
    model = FakeModel(texts=[["Erster Satz", "Zweiter Satz"]])
    run_all(work, model, ["seg1.webm"])
    result = read_result(work, "seg1")
    assert result["v"] == 1
    assert result["captureId"] == "seg1"
    assert result["text"] == "Erster Satz Zweiter Satz"
    assert result["segments"] == [
        {"start": 0.0, "end": 1.5, "text": "Erster Satz"},
        {"start": 2.0, "end": 3.5, "text": "Zweiter Satz"},
    ]
    assert result["language"] == "de"
    assert result["languageProbability"] == 1.0
    assert 9.9 <= result["durationSec"] <= 10.1


def test_passes_the_agreed_whisper_settings(work: Path):
    write_wav(work / "in" / "cfg1.webm", 4)
    model = FakeModel(texts=[["Genug Text hier"]])
    run_all(work, model, ["cfg1.webm"])
    assert model.kwargs[0] == {
        "language": "de",
        "beam_size": 1,
        "vad_filter": True,
        "condition_on_previous_text": False,
        "initial_prompt": "Durchsicht: Koenigin.",
    }


def test_a_text_under_five_characters_is_no_speech(work: Path):
    write_wav(work / "in" / "short1.webm", 4)
    run_all(work, FakeModel(texts=[["Ja"]]), ["short1.webm"])
    assert read_result(work, "short1")["error"] == "NO_SPEECH"


def test_another_language_with_low_probability_is_no_speech(work: Path):
    write_wav(work / "in" / "fr1.webm", 4)
    model = FakeModel(texts=[["Bonjour tout le monde"]], language="fr", probability=0.4)
    run_all(work, model, ["fr1.webm"])
    assert read_result(work, "fr1")["error"] == "NO_SPEECH"


def test_another_language_with_high_probability_keeps_the_text(work: Path):
    write_wav(work / "in" / "en1.webm", 4)
    model = FakeModel(texts=[["Hello there everybody"]], language="en", probability=0.9)
    run_all(work, model, ["en1.webm"])
    assert "error" not in read_result(work, "en1")


def test_no_manifest_means_exit_0_without_loading_the_model(tmp_path: Path, monkeypatch):
    def forbidden(*args, **kwargs):
        raise AssertionError("the model must not load without a manifest")

    monkeypatch.setattr(transcribe, "load_whisper_model", forbidden)
    assert transcribe.main({"VOICE_WORK_DIR": str(tmp_path)}) == 0


def test_an_empty_manifest_does_not_load_the_model(work: Path, monkeypatch):
    (work / "in" / "manifest.json").write_text('{"v":1,"items":[]}', encoding="utf-8")
    monkeypatch.setattr(transcribe, "load_whisper_model", lambda *a: pytest.fail("loaded"))
    assert transcribe.main({"VOICE_WORK_DIR": str(work)}) == 0


def test_a_broken_manifest_fails_the_run(work: Path):
    (work / "in" / "manifest.json").write_text("not json", encoding="utf-8")
    assert transcribe.main({"VOICE_WORK_DIR": str(work)}) == 1


@pytest.mark.parametrize(
    "item",
    [
        {"captureId": "../x", "file": "../x.webm"},
        {"captureId": "a1", "file": "../a1.webm"},
        {"captureId": "a1", "file": "other.webm"},
        {"captureId": "a1", "file": "a1.wav"},
        {"captureId": "", "file": ".webm"},
    ],
)
def test_a_manifest_item_cannot_leave_the_work_dir(work: Path, item: dict):
    manifest = {"v": 1, "items": [item]}
    (work / "in" / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
    with pytest.raises(transcribe.ManifestInvalid):
        transcribe.load_manifest(work)


def test_a_missing_input_file_is_decode_failed(work: Path):
    run_all(work, FakeModel(), ["gone1.webm"])
    assert read_result(work, "gone1")["error"] == "DECODE_FAILED"


@pytest.mark.skipif(os.name == "nt", reason="POSIX file modes")
def test_result_files_are_private(work: Path):
    write_wav(work / "in" / "priv1.webm", 4)
    run_all(work, FakeModel(texts=[["Genug Text hier"]]), ["priv1.webm"])
    mode = stat.S_IMODE((work / "out" / "priv1.json").stat().st_mode)
    assert mode == 0o600


def test_output_never_prints_the_transcript(work: Path, capsys, monkeypatch):
    write_wav(work / "in" / "quiet1.webm", 4)
    make_manifest(work, ["quiet1.webm"])
    model = FakeModel(texts=[["Geheimer Inhalt der Notiz"]])
    vocabulary = work / "vocabulary.txt"
    vocabulary.write_text("Durchsicht: Koenigin.\n", encoding="utf-8")
    monkeypatch.setattr(transcribe, "load_whisper_model", lambda *a: model)
    code = transcribe.main({"VOICE_WORK_DIR": str(work), "WHISPER_VOCABULARY": str(vocabulary)})
    captured = capsys.readouterr()
    assert code == 0
    assert "Geheimer" not in captured.out + captured.err
    assert captured.out.strip() == "transcribed=1 no_speech=0 too_long=0 failed=0"


def test_cpu_threads_are_validated():
    assert transcribe.cpu_threads({}) == 2
    assert transcribe.cpu_threads({"WHISPER_CPU_THREADS": "4"}) == 4
    for bad in ("0", "abc", "99", "-1"):
        with pytest.raises(ValueError):
            transcribe.cpu_threads({"WHISPER_CPU_THREADS": bad})


@pytest.mark.parametrize("extension", ["webm", "ogg", "mp4"])
def test_decodes_each_container_and_hands_the_model_16khz_float_samples(work: Path, extension: str):
    write_wav(work / "in" / f"fmt1.{extension}", 3)
    model = FakeModel(texts=[["Genug Text hier"]])
    run_all(work, model, [f"fmt1.{extension}"])
    assert "error" not in read_result(work, "fmt1")
    audio = model.audio[0]
    assert audio.dtype.name == "float32" and audio.ndim == 1
    assert 2.8 <= audio.size / 16000 <= 3.3


def test_a_file_in_another_container_than_its_extension_is_decode_failed(work: Path):
    write_wav(work / "in" / "ogg1.ogg", 2)
    (work / "in" / "mixed1.webm").write_bytes((work / "in" / "ogg1.ogg").read_bytes())
    model = FakeModel(texts=[["x" * 10]])
    run_all(work, model, ["mixed1.webm"])
    assert read_result(work, "mixed1")["error"] == "DECODE_FAILED"
    assert model.calls == []


def test_audio_without_any_measurable_length_is_decode_failed(work: Path):
    write_wav(work / "in" / "empty1.webm", 0)
    model = FakeModel(texts=[["x" * 10]])
    run_all(work, model, ["empty1.webm"])
    assert read_result(work, "empty1")["error"] == "DECODE_FAILED"
    assert model.calls == []


def test_the_decoded_length_beats_the_demuxed_length(work: Path, monkeypatch):
    write_wav(work / "in" / "liar1.webm", 245)
    monkeypatch.setattr(transcribe, "measure_duration", lambda *a, **k: 10.0)
    model = FakeModel(texts=[["x" * 10]])
    run_all(work, model, ["liar1.webm"])
    assert read_result(work, "liar1")["error"] == "TOO_LONG"
    assert model.calls == []