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
# Own, made-up sample speech; real speeches are only kept locally in testdaten/
SAMPLE_VTT = REPO / "tests" / "daten" / "beispiel.de.vtt"
SAMPLE_TXT = REPO / "tests" / "daten" / "beispiel.de.txt"
SPEECHES = REPO / "testdaten"


@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_vtt_to_lines_without_rolling_window_keeps_identical_cues():
    # Hand-written captions: the speaker says "Nein." twice, as two cues, and
    # a cue may start with the line that ended the previous one
    vtt = "WEBVTT\n\n" + "".join([
        cue(1, "Er sagt:"),
        cue(2, "Nein."),
        cue(3, "Nein."),
        cue(4, "Ja, wirklich", "Nein."),
        cue(5, "Nein.", "Und dann?"),
        cue(6, "Dann kam"),
        cue(7, "der Rest."),
        cue(8, "Ende."),
    ])
    assert vtt2txt.vtt_to_lines(vtt) == [
        "Er sagt:", "Nein.", "Nein.", "Ja, wirklich", "Nein.", "Nein.", "Und dann?",
        "Dann kam", "der Rest.", "Ende.",
    ]


@pytest.mark.parametrize("cues, expected", [
    ([], False),
    ([["A"]], False),
    ([["A"], ["A", "B"]], True),
    ([["A"], ["A", "B"], ["C"]], True),          # half of the pairs overlap
    ([["A"], ["B"], ["C"], ["C", "D"]], False),  # one of three
])
def test_is_rolling(cues, expected):
    assert vtt2txt.is_rolling(cues) is expected


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


def speech(name):
    """A real speech, only kept locally (copyright): skipped without it."""
    path = SPEECHES / name
    return pytest.param(path, marks=pytest.mark.skipif(
        not path.is_file(), reason="echte Reden nur lokal in testdaten/"
    ))


@pytest.mark.parametrize("vtt_path", [
    SAMPLE_VTT,
    speech("ramstein/ramstein.de.vtt"), speech("magdeburg/magdeburg.de-orig.vtt"),
    speech("lahnstein/Lahnstein.de-orig.vtt"),
])
def test_sample_vtt_gives_sample_txt(vtt_path):
    vtt = vtt_path.read_text(encoding="utf-8-sig")
    txt = vtt_path.with_suffix(".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("beispiel.en.vtt"))
    shutil.copy(sample_copy, sample_copy.with_name("anders.de.vtt"))
    run_main(monkeypatch, "-i", sample_copy.parent / "beispiel.*.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 / "beispiel.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/beispiel.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


# === Time format and ranges ===

@pytest.mark.parametrize("text, seconds", [
    ("00:05:10", 310), ("00:20", 1200), ("1:30", 5400), ("25", 1500), ("00:5:00", 300),
    ("00:03:17.5", 197.5), ("00:03:17,5", 197.5), ("1:02:03", 3723),
])
def test_parse_user_time(text, seconds):
    assert vtt2txt.parse_user_time(text) == seconds


@pytest.mark.parametrize("text", ["", "abc", "1:2:3:4", "-5", "1:75", "0:10:60", "1:30.5"])
def test_parse_user_time_rejects_invalid(text):
    with pytest.raises(ValueError):
        vtt2txt.parse_user_time(text)


@pytest.mark.parametrize("text, seconds", [
    ("00:03:17.500", 197.5), ("03:17.500", 197.5), ("1:02:03.000", 3723),
])
def test_parse_cue_time_keeps_minutes_seconds(text, seconds):
    assert vtt2txt.parse_cue_time(text) == seconds


@pytest.mark.parametrize("text, ranges", [
    ("00:5:00-00:6:00 00:10-00:12:00 00:25", [(300, 360), (600, 720), (1500, None)]),
    ("0-00:20", [(None, 1200)]),
    ("-00:20", [(None, 1200)]),
    ("-00:05 00:10-00:12", [(None, 300), (600, 720)]),
    ("0:00:00-0:01", [(None, 60)]),
    ("00:25- 00:10-00:12 0-00:05", [(None, 300), (600, 720), (1500, None)]),
    ("00:05-00:06 00:06-00:07", [(300, 360), (360, 420)]),
])
def test_parse_ranges(text, ranges):
    assert vtt2txt.parse_ranges(text) == ranges


@pytest.mark.parametrize("text, message", [
    ("", "keine Bereiche"),
    ("-", "ungültiger Bereich"),
    ("-0", "Start muss vor Ende"),
    ("-00:05 -00:10", "überschneiden"),
    ("0", "das ganze Video"),
    ("0-", "das ganze Video"),
    ("00:06-00:05", "Start muss vor Ende"),
    ("00:05-00:10 00:08-00:12", "überschneiden"),
    ("00:05 00:10", "überschneiden"),
    ("0-00:05 0-00:10", "überschneiden"),
    ("00:05-x", "ungültiger Bereich"),
])
def test_parse_ranges_rejects(text, message):
    with pytest.raises(ValueError, match=message):
        vtt2txt.parse_ranges(text)


@pytest.mark.parametrize("ranges, suffix", [
    ([], ""),
    ([(197, 325)], "_00-03-17_bis_00-05-25"),
    ([(197, None)], "_ab_00-03-17"),
    ([(None, 325)], "_bis_00-05-25"),
    ([(300, 360), (600, 720), (1500, None)],
     "_00-05-00_bis_00-06-00+00-10-00_bis_00-12-00+ab_00-25-00"),
])
def test_time_suffix(ranges, suffix):
    assert vtt2txt.time_suffix(ranges) == suffix


# === Excerpts ===

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)


def block(vtt, start=None, end=None, exact=False):
    """Lines of the single excerpt of one time range ([] without captions)."""
    blocks = vtt2txt.excerpts(vtt2txt.timed_lines(vtt), [(start, end)], exact)
    return blocks[0][2] if blocks else []


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_exact_time_range():
    vtt = "WEBVTT\n\n" + cue(1, "A") + cue(2, "A", "B") + cue(3, "B", "C") + cue(4, "C", "D")
    assert block(vtt, 2, 4, exact=True) == ["B", "C"]
    assert block(vtt, 3, exact=True) == ["C", "D"]
    assert block(vtt, end=2, exact=True) == ["A"]


def test_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 block(vtt, 2, exact=True) == ["C"]


def test_snap_extends_both_sides_to_sentences():
    # Range holds only "kommt der dritte" and "und endet. Der vierte"
    assert block(SPEECH, 106, 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 block(SPEECH, 102, 104) == ["Der zweite Satz", "beginnt vorher."]
    assert block(SPEECH, 100, 102) == ["Erster Satz ist hier."]


def test_snap_reaches_start_and_end_of_captions():
    assert block(SPEECH, 114) == ["Schluss", "folgt."]
    speech = captions((0, "ohne Satzzeichen"), (2, "am Anfang. Rest"))
    assert block(speech, 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 block(speech, far, far + 1) == ["Satz ohne"]


def test_snap_exact_keeps_cut():
    assert block(SPEECH, 106, 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 block(speech, 2, 4) == ["Und", "dann", "ging er."]


def test_vtt_to_text_headers_with_actual_times():
    text = vtt2txt.vtt_to_text(SPEECH, [(102, 104), (106, 110)])
    assert text == (
        "[00:01:42 - 00:01:45]\nDer zweite Satz\nbeginnt vorher.\n\n"
        "[00:01:44 - 00:01:53]\nDann\nkommt der dritte\nund endet. Der vierte\n"
        "läuft weiter\nund endet auch."
    )


def test_vtt_to_text_merges_blocks_overlapping_after_snap():
    # 106-110 grows to "und endet auch." (cue 112), 110-114 back to "Der vierte" (cue 108)
    text = vtt2txt.vtt_to_text(SPEECH, [(106, 110), (110, 114)])
    assert text == (
        "[00:01:44 - 00:01:55]\nDann\nkommt der dritte\nund endet. Der vierte\n"
        "läuft weiter\nund endet auch. Schluss\nfolgt."
    )


def test_vtt_to_text_shared_line_split_at_sentence_end():
    # Both blocks get their part of "und endet. Der vierte", without overlap
    text = vtt2txt.vtt_to_text(SPEECH, [(106, 108), (110, 112)])
    assert vtt2txt.split_segments(text) == [
        ("[00:01:44 - 00:01:49]", "Dann\nkommt der dritte\nund endet."),
        ("[00:01:48 - 00:01:53]", "Der vierte\nläuft weiter\nund endet auch."),
    ]


def test_vtt_to_text_skips_empty_range(capsys):
    text = vtt2txt.vtt_to_text(SPEECH, [(0, 50), (100, 102), (200, None)])
    assert text == "[00:01:40 - 00:01:41]\nErster Satz ist hier."
    out = capsys.readouterr().out
    assert "keine Untertitel im Bereich 00:00:00 – 00:00:50" in out
    assert "keine Untertitel im Bereich 00:03:20 – Ende (die Untertitel enden bei" in out
    assert vtt2txt.vtt_to_text(SPEECH, [(200, None)]) == ""


def test_vtt_to_text_notes_range_past_last_caption(capsys):
    # "-z 0-05:00" means five hours, not five minutes
    text = vtt2txt.vtt_to_text(SPEECH, [(0, 5 * 3600)])
    assert text.startswith("[00:01:40 - ")
    out = capsys.readouterr().out
    assert "Bereich 00:00:00 – 05:00:00 reicht über das Ende der Untertitel hinaus" in out
    assert "Stunden:Minuten, Minute 5 ist 00:05" in out


@pytest.mark.parametrize("start, end, expected", [
    (100, 110, ""),                           # inside
    (None, None, ""),
    (100, 200, " (die Untertitel enden bei 00:02:00)"),
    (150, None, " (die Untertitel enden bei 00:02:00)"),
    (0, 3600, " (die Untertitel enden bei 00:02:00; zwei Zahlen wie 05:00 sind "
              "Stunden:Minuten, Minute 5 ist 00:05)"),
    (0, 200, ""),                             # no captions at all: no end known
])
def test_end_hint(start, end, expected):
    captions_end = None if (start, end) == (0, 200) else 120
    assert vtt2txt.end_hint(start, end, captions_end) == expected


def test_vtt_to_text_without_ranges_has_no_header():
    assert vtt2txt.vtt_to_text(SPEECH).splitlines()[0] == "Erster Satz ist hier."


def test_split_segments():
    text = "Vorspann\n\n[00:01:00 - 00:02:00]\nA\nB\n\n[applaus]\n[00:03:00 - 00:04:00]\n\nC\n"
    assert vtt2txt.split_segments(text) == [
        (None, "Vorspann"),
        ("[00:01:00 - 00:02:00]", "A\nB\n\n[applaus]"),
        ("[00:03:00 - 00:04:00]", "C"),
    ]
    assert vtt2txt.split_segments("ohne\nÜberschrift") == [(None, "ohne\nÜberschrift")]
    assert vtt2txt.split_segments("  \n") == []


def test_split_segments_inverts_join_segments():
    text = vtt2txt.vtt_to_text(SPEECH, [(102, 104), (112, None)])
    assert vtt2txt.join_segments(vtt2txt.split_segments(text)) == text


# === Command line ===

def test_main_time_range_default_name(sample_copy, monkeypatch):
    run_main(monkeypatch, "-i", sample_copy, "-z", "0:00:30-0:01")
    target = sample_copy.with_name("beispiel.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 vtt2txt.HEADER.fullmatch(part[0])
    assert 0 < len(part) < len(full)


def test_main_several_ranges(sample_copy, monkeypatch):
    run_main(monkeypatch, "-i", sample_copy, "-z", "0:01-0:02 0:05")
    target = sample_copy.with_name("beispiel.de_00-01-00_bis_00-02-00+ab_00-05-00.txt")
    segments = vtt2txt.split_segments(target.read_text(encoding="utf-8"))
    assert [header for header, _ in segments] == [
        "[00:00:59 - 00:02:03]", "[00:04:58 - 00:13:10]",
    ]


@pytest.mark.parametrize("ranges", [["0:01-0:02 0:05"], ["0:01-0:02", "0:05"]])
def test_main_ranges_with_or_without_quotes(sample_copy, monkeypatch, ranges):
    run_main(monkeypatch, "-i", sample_copy, "-z", *ranges)
    target = sample_copy.with_name("beispiel.de_00-01-00_bis_00-02-00+ab_00-05-00.txt")
    assert target.is_file()


@pytest.mark.parametrize("argv, expected", [
    (["-i", "x", "-z", "-00:20"], ["-i", "x", "-z", "0-00:20"]),
    (["-z", "00:05-00:06", "-00:20", "-o", "f"], ["-z", "00:05-00:06", "0-00:20", "-o", "f"]),
    (["--zeitraeume", "-0:01:30"], ["--zeitraeume", "0-0:01:30"]),
    (["-z", "0:05", "--exakt", "-1"], ["-z", "0:05", "--exakt", "-1"]),
    (["-i", "-5.vtt"], ["-i", "-5.vtt"]),
    (["-z=-00:20"], ["-z=-00:20"]),
])
def test_range_argv(argv, expected):
    assert vtt2txt.range_argv(argv) == expected


@pytest.mark.parametrize("ranges", [["-0:01"], ["0-0:01"], ["-z=-0:01"]])
def test_main_range_from_beginning(sample_copy, monkeypatch, ranges):
    args = ranges if ranges[0].startswith("-z") else ["-z", *ranges]
    run_main(monkeypatch, "-i", sample_copy, *args)
    target = sample_copy.with_name("beispiel.de_bis_00-01-00.txt")
    assert target.read_text(encoding="utf-8").startswith("[00:00:00 - ")


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


@pytest.mark.parametrize("args, message", [
    (["-z", "abc"], "ungültiger Bereich"),
    (["-z", "1:00-0:30"], "Start muss vor Ende"),
    (["-z", "0:05 0:10"], "überschneiden"),
    (["-z", "0:05", "0:10"], "überschneiden"),
    (["-z", "-"], "ungültiger Bereich"),
    (["-z", "0"], "das ganze Video"),
    (["-s", "0:01"], "unrecognized arguments: -s"),
])
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 test_main_exact_option(tmp_path, monkeypatch):
    source = tmp_path / "rede.vtt"
    source.write_text(SPEECH, encoding="utf-8")
    run_main(monkeypatch, "-i", source, "-z", "0:01:46-0:01:50", "--exakt")
    target = tmp_path / "rede_00-01-46_bis_00-01-50.txt"
    assert target.read_text(encoding="utf-8") == (
        "[00:01:46 - 00:01:49]\nkommt der dritte\nund endet. Der vierte\n"
    )
    run_main(monkeypatch, "-i", source, "-z", "0:01:46-0:01:50")
    assert target.read_text(encoding="utf-8").endswith("und endet auch.\n")