voice-transcriptor/tests/test_preprocessing.py

import json
from pathlib import Path

import pytest

from voice_transcriptor.models import MediaInfo
from voice_transcriptor.services.preprocessing import CancellationToken, PreprocessingCancelled, PreprocessingOptions, PreprocessingService
from voice_transcriptor.services.job_manifest import JobManifestRepository


class FakeProbe:
    def __init__(self, media: MediaInfo): self.media = media
    def probe(self, path: Path) -> MediaInfo: return self.media


class FakeProcess:
    def __init__(self, command):
        self.stdout = iter(["out_time_us=450000000\n", "progress=continue\n", "progress=end\n"])
        self.returncode = 0
        self.terminated = False
    def wait(self, timeout=None): return self.returncode
    def poll(self): return self.returncode
    def terminate(self): self.terminated = True
    def kill(self): self.terminated = True


def make_media(path: Path, duration="1800") -> MediaInfo:
    path.write_bytes(b"source")
    return MediaInfo(path, 6, float(duration), "aac", duration, True, path.suffix in {".mp4", ".mov", ".mkv", ".webm"})


def test_preprocess_builds_streaming_commands_and_exact_manifest(tmp_path: Path) -> None:
    source = tmp_path / "recording.mp4"
    commands = []
    def factory(command, **kwargs):
        commands.append(command); Path(command[-1]).write_bytes(b"chunk"); return FakeProcess(command)
    progress = []
    service = PreprocessingService(Path("ffmpeg"), FakeProbe(make_media(source)), tmp_path / "jobs", factory)
    result = service.preprocess(source, PreprocessingOptions(), progress=progress.append)
    manifest = json.loads(result.manifest_path.read_text(encoding="utf-8"))
    assert len(commands) == 3
    assert [commands[0][commands[0].index(f) + 1] for f in ("-ss", "-t", "-ac", "-ar", "-c:a", "-b:a")] == ["0", "900", "1", "24000", "aac", "64k"]
    assert all("0:a:0" in command and "-progress" in command for command in commands)
    assert manifest["state"] == "completed"
    assert manifest["chunks"][1]["source_start_seconds"] == "885"
    assert manifest["chunks"][-1]["source_end_seconds"] == "1800"
    assert manifest["completed_chunks"] == 3
    assert progress[-1].percent == 100
    assert [p.percent for p in progress] == sorted(p.percent for p in progress)


def test_cleanup_and_retention_are_job_scoped(tmp_path: Path) -> None:
    source = tmp_path / "recording.wav"
    def factory(command, **kwargs): Path(command[-1]).write_bytes(b"x"); return FakeProcess(command)
    service = PreprocessingService(Path("ffmpeg"), FakeProbe(make_media(source, "10")), tmp_path / "jobs", factory)
    result = service.preprocess(source, PreprocessingOptions()); job = result.job_directory; result.cleanup(); assert not job.exists()
    retained = service.preprocess(source, PreprocessingOptions(retain_temporary_files=True)); retained.cleanup(); assert retained.job_directory.exists()


def test_durable_preprocessing_creates_schema_2_pending_transcription_job(tmp_path: Path) -> None:
    source = tmp_path / "recording.wav"
    output = tmp_path / "output"

    def factory(command, **kwargs):
        Path(command[-1]).write_bytes(b"encoded")
        return FakeProcess(command)

    service = PreprocessingService(
        Path("ffmpeg"),
        FakeProbe(make_media(source, "10")),
        process_factory=factory,
        manifest_repository=JobManifestRepository(),
    )

    result = service.preprocess(source, PreprocessingOptions(), durable_root=output)
    manifest = json.loads(result.manifest_path.read_text(encoding="utf-8"))

    assert result.job_directory.parent == (output / "voice-transcriptor-jobs").resolve()
    assert result.retained is True
    assert manifest["schema_version"] == 2
    assert manifest["state"] == "transcribing"
    assert manifest["completed_chunks"] == 0
    assert manifest["chunks"][0]["status"] == "pending"
    assert manifest["chunks"][0]["encoded"] is True
    result.cleanup()
    assert result.job_directory.exists()


def test_cancelled_job_stops_before_process(tmp_path: Path) -> None:
    source = tmp_path / "recording.m4a"; token = CancellationToken(); token.cancel()
    service = PreprocessingService(Path("ffmpeg"), FakeProbe(make_media(source, "10")), tmp_path / "jobs", lambda *a, **k: pytest.fail("started"))
    with pytest.raises(PreprocessingCancelled): service.preprocess(source, PreprocessingOptions(), token)


@pytest.mark.parametrize("suffix", [".m4a", ".mp3", ".wav", ".mp4", ".mov", ".webm", ".mkv"])
def test_supported_extensions(tmp_path: Path, suffix: str) -> None:
    source = tmp_path / f"recording{suffix}"
    def factory(command, **kwargs): Path(command[-1]).write_bytes(b"x"); return FakeProcess(command)
    result = PreprocessingService(Path("ffmpeg"), FakeProbe(make_media(source, "1")), tmp_path / "jobs", factory).preprocess(source, PreprocessingOptions())
    assert result.manifest_path.exists()