tests/test_longform.py

from __future__ import annotations

import json
import wave
from pathlib import Path

import pytest

from src.jobs import JobPaths
from src.longform import (
    CHUNK_MODE_AUTO,
    CHUNK_MODE_FIXED,
    RANGE_MODE_CUSTOM,
    RANGE_MODE_PREVIEW,
    build_long_form_plan,
)
from src.media import prepare_inputs, probe_duration


def test_preview_range_is_clamped_per_input():
    plan = build_long_form_plan(
        [100.0, 50.0],
        range_mode=RANGE_MODE_PREVIEW,
        range_start_seconds=40,
        preview_seconds=30,
        chunk_mode=CHUNK_MODE_AUTO,
    )
    assert plan.total_selected_seconds == 40.0
    assert plan.input_ranges[0]["end_seconds"] == 70.0
    assert plan.input_ranges[1]["end_seconds"] == 50.0
    assert plan.input_ranges[1]["source_end_clamped"] is True
    assert plan.resolved_chunk_seconds == 0


def test_custom_range_rejects_empty_selection():
    with pytest.raises(ValueError, match="outside input"):
        build_long_form_plan(
            [20.0],
            range_mode=RANGE_MODE_CUSTOM,
            range_start_seconds=20,
            range_end_seconds=25,
        )


def test_auto_chunk_policy_uses_5_or_10_minute_chunks():
    medium = build_long_form_plan([1800], chunk_mode=CHUNK_MODE_AUTO, model_count=1)
    long_single = build_long_form_plan([7200], chunk_mode=CHUNK_MODE_AUTO, model_count=1)
    long_ensemble = build_long_form_plan([7200], chunk_mode=CHUNK_MODE_AUTO, model_count=2)
    assert medium.resolved_chunk_seconds == 300
    assert long_single.resolved_chunk_seconds == 600
    assert long_ensemble.resolved_chunk_seconds == 300
    assert long_single.chunk_count_total == 12


def test_fixed_chunk_rejects_unknown_duration():
    with pytest.raises(ValueError, match="120, 300, or 600"):
        build_long_form_plan([3600], chunk_mode=CHUNK_MODE_FIXED, fixed_chunk_seconds=240)


def _write_silence(path: Path, seconds: int = 4) -> None:
    with wave.open(str(path), "wb") as handle:
        handle.setnchannels(2)
        handle.setsampwidth(2)
        handle.setframerate(44100)
        handle.writeframes(b"\x00\x00\x00\x00" * 44100 * seconds)


def test_prepare_inputs_extracts_selected_range_before_gpu(tmp_path):
    source = tmp_path / "source.wav"
    _write_silence(source, 4)
    root = tmp_path / "job"
    job = JobPaths("a" * 32, root, root / "input", root / "work", root / "output")
    for directory in (job.input_dir, job.work_dir, job.output_dir, job.logs_dir, job.bundle_dir, job.config_dir):
        directory.mkdir(parents=True, exist_ok=True)

    prepared, total = prepare_inputs(
        job,
        [source],
        input_ranges=[{"start_seconds": 1.0, "end_seconds": 3.0}],
    )
    assert total == pytest.approx(2.0, abs=0.01)
    assert prepared[0].parent == job.work_dir
    assert prepared[0].suffix == ".wav"
    assert probe_duration(prepared[0]) == pytest.approx(2.0, abs=0.03)


def test_ui_wrapper_maps_long_form_fields(monkeypatch):
    from src import ui

    captured = {}

    def fake_prepare(*args, **kwargs):
        captured["args"] = args
        captured["kwargs"] = kwargs
        return "state", "markdown", "preflight.json", "prepare.log"

    monkeypatch.setattr(ui, "prepare_job", fake_prepare)
    values = list(range(34))
    result = ui.prepare_separation_with_progress(*values, progress=None)
    assert result[:4] == ("state", "markdown", "preflight.json", "prepare.log")
    assert len(captured["args"]) == 29
    assert captured["args"][22] == 0
    assert captured["kwargs"]["range_mode"] == 22
    assert captured["kwargs"]["range_start_seconds"] == 23
    assert captured["kwargs"]["range_end_seconds"] == 24
    assert captured["kwargs"]["preview_seconds"] == 25
    assert captured["kwargs"]["chunk_mode"] == 26
    assert captured["kwargs"]["fixed_chunk_seconds"] == 27
    assert captured["kwargs"]["batch_continue_on_item_error"] is True