second-brain/tests/test_transcribe.py

"""Unit + integration coverage for second_brain.transcribe.

Unit-level: resolve_settings + write_srt are pure-Python and run anywhere.

Integration: an end-to-end CPU/int8 transcribe of a synthesized audio
clip — auto-skipped when ffmpeg or faster-whisper is missing (the dev
side doesn't install faster-whisper by default). Only the *tower* with
GPU large-v3 can validate the production hot path; this test proves the
wiring + the SRT contract + the lazy-import path on cheap CPU.
"""

from __future__ import annotations


import pytest


# ---------------------------------------------------------------------------
# Pure helpers
# ---------------------------------------------------------------------------


def test_resolve_settings_picks_int8_for_cpu_by_default():
    from second_brain.transcribe import resolve_settings

    out = resolve_settings({"whisper_device": "cpu"})
    assert out["device"] == "cpu"
    assert out["compute_type"] == "int8"
    assert out["model"] == "large-v3"


def test_resolve_settings_picks_float16_for_cuda_by_default():
    from second_brain.transcribe import resolve_settings

    out = resolve_settings({"whisper_device": "cuda"})
    assert out["device"] == "cuda"
    assert out["compute_type"] == "float16"


def test_resolve_settings_env_overrides_block(monkeypatch):
    from second_brain.transcribe import resolve_settings

    monkeypatch.setenv("WHISPER_MODEL", "small")
    monkeypatch.setenv("WHISPER_DEVICE", "cpu")
    monkeypatch.setenv("WHISPER_COMPUTE_TYPE", "int8_float16")
    out = resolve_settings({"whisper_device": "cuda"})
    assert out["model"] == "small"
    assert out["device"] == "cpu"
    assert out["compute_type"] == "int8_float16"


def test_write_srt_roundtrip(tmp_path):
    from second_brain.transcribe import write_srt

    segs = [
        {"start": 0.0, "end": 1.5, "text": "hello"},
        {"start": 1.5, "end": 3.25, "text": "world"},
    ]
    out = tmp_path / "x.srt"
    write_srt(segs, out)
    body = out.read_text()
    assert "1\n00:00:00,000 --> 00:00:01,500\nhello" in body
    assert "2\n00:00:01,500 --> 00:00:03,250\nworld" in body


# ---------------------------------------------------------------------------
# Segment / word conversion — exercised on faked faster-whisper objects so
# we don't need a real model + audio to validate the JSONB-bound shape.
# ---------------------------------------------------------------------------


class _FakeWord:
    def __init__(self, start, end, word, probability):
        self.start = start
        self.end = end
        self.word = word
        self.probability = probability


class _FakeSegment:
    def __init__(self, id, start, end, text, words=None):
        self.id = id
        self.start = start
        self.end = end
        self.text = text
        self.words = words


def test_segment_to_dict_with_words():
    """Happy path: faster-whisper segment + word_timestamps → JSON-safe dict."""
    import json

    from second_brain.transcribe import segment_to_dict

    seg = _FakeSegment(
        id=3,
        start=1.25,
        end=3.5,
        text=" hello world  ",
        words=[
            _FakeWord(1.25, 1.6, " hello", 0.92),
            _FakeWord(1.6, 3.5, " world", 0.87),
        ],
    )
    d = segment_to_dict(seg)

    assert d == {
        "id": 3,
        "start": 1.25,
        "end": 3.5,
        "text": "hello world",  # stripped
        "words": [
            {"start": 1.25, "end": 1.6, "word": " hello", "probability": 0.92},
            {"start": 1.6, "end": 3.5, "word": " world", "probability": 0.87},
        ],
    }
    # Must be JSON-serialisable for the JSONB column.
    json.dumps(d)


def test_segment_to_dict_without_words():
    """word_timestamps=False (or older faster-whisper) → segment.words is None."""
    from second_brain.transcribe import segment_to_dict

    seg = _FakeSegment(id=0, start=0.0, end=1.0, text="just text", words=None)
    d = segment_to_dict(seg)
    assert d["text"] == "just text"
    assert d["words"] == []


def test_segment_to_dict_skips_malformed_words():
    """A single bad word entry must not drop the whole segment."""
    from second_brain.transcribe import segment_to_dict

    class _BadWord:
        # Missing start → float(None) raises in the helper, must be skipped.
        start = None
        end = 1.0
        word = "broken"
        probability = 0.5

    seg = _FakeSegment(
        id=1,
        start=0.0,
        end=1.0,
        text="hi",
        words=[_FakeWord(0.0, 0.5, "hi", 0.99), _BadWord()],
    )
    d = segment_to_dict(seg)
    # Good word survives; bad word is filtered.
    assert len(d["words"]) == 1
    assert d["words"][0]["word"] == "hi"


# ---------------------------------------------------------------------------
# Integration: synthesized audio → CPU/int8 faster-whisper
# ---------------------------------------------------------------------------


def _have_faster_whisper() -> bool:
    try:
        import faster_whisper  # type: ignore[import-not-found]  # noqa: F401
    except ImportError:
        return False
    return True


def _have_numpy() -> bool:
    try:
        import numpy  # noqa: F401
    except ImportError:
        return False
    return True


@pytest.mark.skipif(
    not (_have_faster_whisper() and _have_numpy()),
    reason="needs faster-whisper (uv sync --extra tower) + numpy to exercise CPU path",
)
def test_cpu_transcribe_smoke(tmp_path):
    """Hand faster-whisper a synthesized audio array on CPU/int8 with `tiny`.

    Pure-tone input gives empty / minimal recognition — the point is the
    wiring: model loads, the segments iterator drains, an SRT is writable.
    Using `tiny` to keep this under ~30s coldcache. Skipped on dev where
    faster-whisper isn't installed; only the tower's `uv sync --extra
    tower` brings it in.

    Bypassing ffmpeg by passing a numpy array — faster-whisper.transcribe
    accepts (path | file-like | np.ndarray). This exercises the same code
    path the worker uses, only with zero external binaries.
    """
    import numpy as np

    from second_brain.transcribe import write_srt

    # Build 2s of 440Hz mono at 16kHz, the model's native sample rate.
    sr = 16000
    duration = 2.0
    t = np.arange(0, int(sr * duration)) / sr
    audio = (0.2 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)

    # Load `tiny` to keep this cheap; we're not asserting transcription
    # quality, just that the pipeline executes and returns the documented
    # shape.
    from faster_whisper import WhisperModel  # type: ignore[import-not-found]

    model = WhisperModel("tiny", device="cpu", compute_type="int8")
    segments_iter, _ = model.transcribe(audio, language="en", beam_size=1)

    # Materialize to mirror what Transcriber.transcribe does internally.
    segments: list[dict] = []
    text_parts: list[str] = []
    for seg in segments_iter:
        s = (seg.text or "").strip()
        segments.append({"start": float(seg.start), "end": float(seg.end), "text": s})
        if s:
            text_parts.append(s)
    text = " ".join(text_parts)

    assert isinstance(text, str)
    assert isinstance(segments, list)

    srt = tmp_path / "out.srt"
    write_srt(segments, srt)
    assert srt.exists()