second-brain/tests/test_transcribe.py
Travis Herbranson 913d4f0335 deploy/tower + claim race + transcribe smoke tests
deploy/tower/:
- second-brain-transcribe.service — systemd unit. User=herbyadmin,
  Type=simple, After=/Wants= wg-quick@wg-lan.service so the WG tunnel
  must come up first. Restart=always with a StartLimitBurst guard.
- second-brain-transcribe.env.example — env file template documenting
  the SECOND_BRAIN_DATABASE_URL form for db.wg.herbylab.dev (10.99.0.1)
  and the optional WHISPER_* overrides.
- README.md — EndeavourOS install steps (nvidia/cuda/cudnn, ffmpeg, uv
  + tower extra, model pre-warm), WG topology reference, validation
  checklist for what to confirm once the tunnel is live, and a
  follow-ups section flagging the local-disk → NAS media migration as
  out-of-scope-for-this-round.

Tests:
- tests/test_claim.py — live-DB race test. Two threads call
  claim_next_source against a single PULLED video row; SKIP LOCKED
  must give exactly one of them the row, the other gets None. Also
  asserts the claimed_by/at columns land + release nulls them.
  Auto-skips when no SECOND_BRAIN_DATABASE_URL is set.
- tests/test_transcribe.py — pure-Python coverage of resolve_settings
  (cpu→int8, cuda→float16, env-over-block) and write_srt; plus a CPU
  smoke test that synthesizes a numpy audio array and runs the `tiny`
  model on cpu/int8 (auto-skipped when faster-whisper isn't installed,
  i.e. on the dev side without --extra tower).
2026-05-25 10:20:46 -04:00

143 lines
4.7 KiB
Python

"""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 os
import shutil
import subprocess
import tempfile
from pathlib import Path
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
# ---------------------------------------------------------------------------
# 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()