youtube_abschrift_korrigieren/tests/test_vtt2txt.py

"""Tests for vtt2txt.py."""
import shutil
import sys
from pathlib import Path

import pytest

import vtt2txt

REPO = Path(__file__).resolve().parent.parent
SAMPLE_VTT = REPO / "testdaten" / "ramstein.de.vtt"
SAMPLE_TXT = REPO / "testdaten" / "ramstein.de.txt"


@pytest.fixture
def sample_copy(tmp_path):
    target = tmp_path / SAMPLE_VTT.name
    shutil.copy(SAMPLE_VTT, target)
    return target


def cue(start, *lines):
    return f"00:00:{start:02d}.000 --> 00:00:{start + 1:02d}.000\n" + "\n".join(lines) + "\n\n"


def run_main(monkeypatch, *args):
    monkeypatch.setattr(sys, "argv", ["vtt2txt.py", *map(str, args)])
    vtt2txt.main()


def test_vtt_to_lines_strips_tags_and_rolling_duplicates():
    vtt = (
        "WEBVTT\nKind: captions\n\n"
        "00:00:01.000 --> 00:00:02.000 align:start\n"
        "Erste<00:00:01.500><c> Zeile</c>\n\n"
        "00:00:02.000 --> 00:00:03.000\nErste Zeile\n \n\n"
        "cue-id\n00:00:03.000 --> 00:00:04.000\nErste Zeile\nA &amp; B\n"
    )
    assert vtt2txt.vtt_to_lines(vtt) == ["Erste Zeile", "A & B"]


def test_vtt_to_lines_keeps_text_after_space_only_line():
    # YouTube cues start with a line holding a single space; the text after it
    # must survive even if no later cue repeats it (here: end of file)
    vtt = "WEBVTT\n\n00:00:01.000 --> 00:00:02.000\n \n[applaus]\n"
    assert vtt2txt.vtt_to_lines(vtt) == ["[applaus]"]


def test_vtt_to_lines_skips_note_and_normalizes_whitespace():
    vtt = (
        "WEBVTT\n\nNOTE Kommentar\nnicht übernehmen\n\n"
        "00:00:01.000 --> 00:00:02.000\nA&nbsp; B\n\n"
        "00:00:02.000 --> 00:00:03.000\nA B\nC\n"
    )
    assert vtt2txt.vtt_to_lines(vtt) == ["A B", "C"]


def test_vtt_to_lines_keeps_real_repetition():
    # Rolling window: each cue repeats the previous line, then adds new text.
    # The speaker really says "Nein." twice; both must survive.
    vtt = "WEBVTT\n\n" + "".join([
        cue(1, " ", "Er sagt:"),
        cue(2, "Er sagt:", " "),
        cue(3, "Er sagt:", "Nein."),
        cue(4, "Nein.", " "),
        cue(5, "Nein.", "Nein."),
        cue(6, " ", " "),
        cue(7, "Nein.", "Ende."),
    ])
    assert vtt2txt.vtt_to_lines(vtt) == ["Er sagt:", "Nein.", "Nein.", "Ende."]


def test_vtt_to_lines_without_rolling_window():
    vtt = "WEBVTT\n\n1\n" + cue(1, "A", "B") + "2\n" + cue(2, "C")
    assert vtt2txt.vtt_to_lines(vtt) == ["A", "B", "C"]


def test_main_default_output_matches_sample(sample_copy, monkeypatch, capsys):
    run_main(monkeypatch, "-i", sample_copy)
    target = sample_copy.with_name("ramstein.de.txt")
    assert target.read_text(encoding="utf-8") == SAMPLE_TXT.read_text(encoding="utf-8")
    assert f"Erstellt: {target}" in capsys.readouterr().out


@pytest.mark.parametrize("name", ["ramstein.de", "magdeburg.de-orig", "Lahnstein.de-orig"])
def test_sample_vtt_gives_sample_txt(name):
    vtt = (REPO / "testdaten" / f"{name}.vtt").read_text(encoding="utf-8-sig")
    txt = (REPO / "testdaten" / f"{name}.txt").read_text(encoding="utf-8")
    assert vtt2txt.vtt_to_lines(vtt) == txt.splitlines()


def test_main_multiple_inputs(sample_copy, monkeypatch):
    second = sample_copy.with_name("rede.en.vtt")
    shutil.copy(sample_copy, second)
    run_main(monkeypatch, "-i", sample_copy, second)
    assert sample_copy.with_suffix(".txt").is_file()
    assert second.with_suffix(".txt").is_file()


def test_main_expands_glob_without_duplicates(sample_copy, monkeypatch, capsys):
    shutil.copy(sample_copy, sample_copy.with_name("ramstein.en.vtt"))
    shutil.copy(sample_copy, sample_copy.with_name("anders.de.vtt"))
    run_main(monkeypatch, "-i", sample_copy.parent / "ramstein.*.vtt", sample_copy)
    out = capsys.readouterr().out
    assert out.count("Erstellt:") == 2
    assert not sample_copy.with_name("anders.de.txt").exists()


@pytest.mark.parametrize("make_dir, suffix", [(True, ""), (False, "/")])
def test_main_output_directory(sample_copy, monkeypatch, make_dir, suffix):
    second = sample_copy.with_name("rede.en.vtt")
    shutil.copy(sample_copy, second)
    out_dir = sample_copy.parent / "out"
    if make_dir:
        out_dir.mkdir()
    run_main(monkeypatch, "-i", sample_copy, second, "-o", f"{out_dir}{suffix}")
    assert (out_dir / "ramstein.de.txt").is_file()
    assert (out_dir / "rede.en.txt").is_file()


def test_main_output_option_creates_directories(sample_copy, monkeypatch):
    target = sample_copy.parent / "neu" / "unter" / "rede.txt"
    run_main(monkeypatch, "-i", sample_copy, "-o", target)
    assert target.is_file()


@pytest.mark.parametrize("content, message", [
    ("WEBVTT\n".encode("utf-8"), "Keine Untertitel"),
    ("WEBVTT\n\n00:00:01.000 --> 00:00:02.000\nGrüße\n".encode("latin-1"),
     "Keine UTF-8-Datei"),
])
def test_main_bad_file_continues_batch(sample_copy, monkeypatch, capsys, content, message):
    bad = sample_copy.with_name("aaa.vtt")
    bad.write_bytes(content)
    with pytest.raises(SystemExit) as exc:
        run_main(monkeypatch, "-i", bad, sample_copy)
    assert exc.value.code == 1
    assert message in capsys.readouterr().err
    assert not bad.with_suffix(".txt").exists()
    assert sample_copy.with_suffix(".txt").is_file()


@pytest.mark.parametrize("args, message", [
    ([], "required"),
    (["-i", "gibtsnicht.vtt"], "Eingabedatei nicht gefunden"),
    (["-i", "{txt}"], "Keine VTT-Datei"),
    (["-i", "{vtt}", "{vtt2}", "-o", "{dir}/x.txt"], "muss -o/--output ein Verzeichnis sein"),
    (["-i", "{dir}/*.gibtsnicht"], "Keine Datei passt zu"),
    (["-i", "{vtt}", "{dir}/sub/ramstein.de.vtt", "-o", "{dir}/"], "dieselbe Ausgabedatei"),
    (["-i", "{vtt}", "-o", "{vtt}"], "nicht die Eingabedatei"),
])
def test_cli_errors(sample_copy, monkeypatch, capsys, args, message):
    txt = sample_copy.with_suffix(".txt")
    txt.write_text("x\n", encoding="utf-8")
    vtt2 = sample_copy.with_name("rede.en.vtt")
    shutil.copy(sample_copy, vtt2)
    (sample_copy.parent / "sub").mkdir()
    shutil.copy(sample_copy, sample_copy.parent / "sub" / sample_copy.name)
    args = [a.format(vtt=sample_copy, vtt2=vtt2, txt=txt, dir=sample_copy.parent)
            for a in args]
    with pytest.raises(SystemExit) as exc:
        run_main(monkeypatch, *args)
    assert exc.value.code == 2
    assert message in capsys.readouterr().err


@pytest.mark.parametrize("text, seconds", [
    ("00:03:17", 197), ("3:17", 197), ("197", 197), ("00:03:17.5", 197.5),
    ("00:03:17,5", 197.5), ("1:02:03", 3723),
])
def test_parse_time(text, seconds):
    assert vtt2txt.parse_time(text) == seconds


@pytest.mark.parametrize("text", ["", "abc", "1:2:3:4", "-5"])
def test_parse_time_rejects_invalid(text):
    with pytest.raises(ValueError):
        vtt2txt.parse_time(text)


@pytest.mark.parametrize("start, end, suffix", [
    (None, None, ""),
    (197, 325, "_00-03-17_bis_00-05-25"),
    (197, None, "_ab_00-03-17"),
    (None, 325, "_bis_00-05-25"),
])
def test_time_suffix(start, end, suffix):
    assert vtt2txt.time_suffix(start, end) == suffix


def test_vtt_to_lines_exact_time_range():
    vtt = "WEBVTT\n\n" + cue(1, "A") + cue(2, "A", "B") + cue(3, "B", "C") + cue(4, "C", "D")
    assert vtt2txt.vtt_to_lines(vtt, start=2, end=4, exact=True) == ["B", "C"]
    assert vtt2txt.vtt_to_lines(vtt, start=3, exact=True) == ["C", "D"]
    assert vtt2txt.vtt_to_lines(vtt, end=2, exact=True) == ["A"]


def test_vtt_to_lines_time_range_skips_window_from_before_start():
    # "B" was said before the range; its repetition in the first cue must not count
    vtt = "WEBVTT\n\n" + cue(1, "A", "B") + cue(2, "B", "C")
    assert vtt2txt.vtt_to_lines(vtt, start=2, exact=True) == ["C"]


def test_main_time_range_default_name(sample_copy, monkeypatch):
    run_main(monkeypatch, "-i", sample_copy, "-s", "0:30", "-e", "1:00")
    target = sample_copy.with_name("ramstein.de_00-00-30_bis_00-01-00.txt")
    full = SAMPLE_TXT.read_text(encoding="utf-8").splitlines()
    part = target.read_text(encoding="utf-8").splitlines()
    assert 0 < len(part) < len(full)


def test_main_time_range_without_captions(sample_copy, monkeypatch, capsys):
    with pytest.raises(SystemExit) as exc:
        run_main(monkeypatch, "-i", sample_copy, "-s", "99:00:00")
    assert exc.value.code == 1
    assert "im gewählten Zeitraum" in capsys.readouterr().err


@pytest.mark.parametrize("args, message", [
    (["-s", "abc"], "ungültige Zeit"),
    (["-s", "1:00", "-e", "0:30"], "muss vor -e/--ende liegen"),
])
def test_cli_time_errors(sample_copy, monkeypatch, capsys, args, message):
    with pytest.raises(SystemExit) as exc:
        run_main(monkeypatch, "-i", sample_copy, *args)
    assert exc.value.code == 2
    assert message in capsys.readouterr().err


def timed_cue(seconds, *lines):
    return f"{vtt2txt.format_time(seconds)}.000 --> {vtt2txt.format_time(seconds + 1)}.000\n" \
        + "\n".join(lines) + "\n\n"


def captions(*cues):
    """VTT with one new line per cue: cues are (seconds, text)."""
    return "WEBVTT\n\n" + "".join(timed_cue(t, text) for t, text in cues)


SPEECH = captions(
    (100, "Erster Satz ist hier."), (102, "Der zweite Satz"), (104, "beginnt vorher. Dann"),
    (106, "kommt der dritte"), (108, "und endet. Der vierte"), (110, "läuft weiter"),
    (112, "und endet auch. Schluss"), (114, "folgt."),
)


def test_snap_extends_both_sides_to_sentences():
    # Range holds only "kommt der dritte" and "und endet. Der vierte"
    assert vtt2txt.vtt_to_lines(SPEECH, start=106, end=110) == [
        "Dann", "kommt der dritte", "und endet. Der vierte", "läuft weiter",
        "und endet auch.",
    ]


def test_snap_keeps_range_already_at_sentence_bounds():
    assert vtt2txt.vtt_to_lines(SPEECH, start=102, end=104) == [
        "Der zweite Satz", "beginnt vorher.",
    ]
    assert vtt2txt.vtt_to_lines(SPEECH, start=100, end=102) == ["Erster Satz ist hier."]


def test_snap_reaches_start_and_end_of_captions():
    assert vtt2txt.vtt_to_lines(SPEECH, start=114) == ["Schluss", "folgt."]
    speech = captions((0, "ohne Satzzeichen"), (2, "am Anfang. Rest"))
    assert vtt2txt.vtt_to_lines(speech, start=2) == ["ohne Satzzeichen", "am Anfang. Rest"]


def test_snap_cuts_exactly_without_sentence_end_in_reach():
    far = vtt2txt.MAX_SNAP_SECONDS + 5
    speech = captions(
        (0, "Satz."), (2, "ein endloser"), (far, "Satz ohne"), (2 * far, "Ende"), (3 * far, "jetzt."),
    )
    assert vtt2txt.vtt_to_lines(speech, start=far, end=far + 1) == ["Satz ohne"]


def test_snap_exact_keeps_cut():
    assert vtt2txt.vtt_to_lines(SPEECH, start=106, end=110, exact=True) == [
        "kommt der dritte", "und endet. Der vierte",
    ]


def test_snap_closing_quote_belongs_to_sentence():
    speech = captions((0, "Er sagte „Nein.“ Und"), (2, "dann"), (4, "ging er."))
    assert vtt2txt.vtt_to_lines(speech, start=2, end=4) == ["Und", "dann", "ging er."]


def test_main_exact_option(tmp_path, monkeypatch):
    source = tmp_path / "rede.vtt"
    source.write_text(SPEECH, encoding="utf-8")
    run_main(monkeypatch, "-i", source, "-s", "106", "-e", "110", "--exakt")
    target = tmp_path / "rede_00-01-46_bis_00-01-50.txt"
    assert target.read_text(encoding="utf-8") == "kommt der dritte\nund endet. Der vierte\n"
    run_main(monkeypatch, "-i", source, "-s", "106", "-e", "110")
    assert target.read_text(encoding="utf-8").endswith("und endet auch.\n")