chrysopedia/backend/pipeline/test_caption_generator.py

"""Unit tests for caption_generator module."""

from __future__ import annotations

import re
import tempfile
from pathlib import Path

import pytest

from pipeline.caption_generator import (
    DEFAULT_STYLE,
    _format_ass_time,
    generate_ass_captions,
    write_ass_file,
)


# ── Fixtures ─────────────────────────────────────────────────────────────────

@pytest.fixture
def sample_word_timings() -> list[dict]:
    """Realistic word timings as produced by extract_word_timings."""
    return [
        {"word": "This", "start": 10.0, "end": 10.3},
        {"word": "is", "start": 10.3, "end": 10.5},
        {"word": "a", "start": 10.5, "end": 10.6},
        {"word": "test", "start": 10.6, "end": 11.0},
        {"word": "sentence", "start": 11.1, "end": 11.6},
    ]


# ── Time formatting ─────────────────────────────────────────────────────────

class TestFormatAssTime:
    def test_zero(self):
        assert _format_ass_time(0.0) == "0:00:00.00"

    def test_sub_second(self):
        assert _format_ass_time(0.5) == "0:00:00.50"

    def test_minutes(self):
        assert _format_ass_time(65.5) == "0:01:05.50"

    def test_hours(self):
        assert _format_ass_time(3661.25) == "1:01:01.25"

    def test_negative_clamps_to_zero(self):
        assert _format_ass_time(-5.0) == "0:00:00.00"


# ── ASS generation ──────────────────────────────────────────────────────────

class TestGenerateAssCaptions:
    def test_empty_timings_returns_header_only(self):
        result = generate_ass_captions([], clip_start=0.0)
        assert "[Script Info]" in result
        assert "[Events]" in result
        # No Dialogue lines
        assert "Dialogue:" not in result

    def test_structure_has_required_sections(self, sample_word_timings):
        result = generate_ass_captions(sample_word_timings, clip_start=10.0)
        assert "[Script Info]" in result
        assert "[V4+ Styles]" in result
        assert "[Events]" in result
        assert "Dialogue:" in result

    def test_clip_offset_applied(self, sample_word_timings):
        """Word at t=10.5 with clip_start=10.0 should become t=0.5 in ASS."""
        result = generate_ass_captions(sample_word_timings, clip_start=10.0)
        lines = result.strip().split("\n")
        dialogue_lines = [l for l in lines if l.startswith("Dialogue:")]

        # First word "This" starts at 10.0, clip_start=10.0 → relative 0.0
        assert dialogue_lines[0].startswith("Dialogue: 0,0:00:00.00,")

        # Third word "a" starts at 10.5, clip_start=10.0 → relative 0.5
        assert "0:00:00.50" in dialogue_lines[2]

    def test_karaoke_tags_present(self, sample_word_timings):
        result = generate_ass_captions(sample_word_timings, clip_start=10.0)
        lines = result.strip().split("\n")
        dialogue_lines = [l for l in lines if l.startswith("Dialogue:")]

        for line in dialogue_lines:
            # Each line should have a \kN tag
            assert re.search(r"\{\\k\d+\}", line), f"Missing karaoke tag in: {line}"

    def test_karaoke_duration_math(self, sample_word_timings):
        """Word "This" at [10.0, 10.3] → 0.3s → k30 (30 centiseconds)."""
        result = generate_ass_captions(sample_word_timings, clip_start=10.0)
        lines = result.strip().split("\n")
        dialogue_lines = [l for l in lines if l.startswith("Dialogue:")]

        # "This" duration: 10.3 - 10.0 = 0.3s = 30cs
        assert "{\\k30}This" in dialogue_lines[0]

        # "test" duration: 11.0 - 10.6 = 0.4s = 40cs
        assert "{\\k40}test" in dialogue_lines[3]

    def test_word_count_matches(self, sample_word_timings):
        result = generate_ass_captions(sample_word_timings, clip_start=10.0)
        lines = result.strip().split("\n")
        dialogue_lines = [l for l in lines if l.startswith("Dialogue:")]
        assert len(dialogue_lines) == 5

    def test_empty_word_text_skipped(self):
        timings = [
            {"word": "hello", "start": 0.0, "end": 0.5},
            {"word": "  ", "start": 0.5, "end": 0.7},  # whitespace-only
            {"word": "", "start": 0.7, "end": 0.8},     # empty
            {"word": "world", "start": 0.8, "end": 1.2},
        ]
        result = generate_ass_captions(timings, clip_start=0.0)
        lines = result.strip().split("\n")
        dialogue_lines = [l for l in lines if l.startswith("Dialogue:")]
        assert len(dialogue_lines) == 2  # only "hello" and "world"

    def test_custom_style_overrides(self, sample_word_timings):
        result = generate_ass_captions(
            sample_word_timings,
            clip_start=10.0,
            style_config={"font_size": 72, "font_name": "Roboto"},
        )
        assert "Roboto" in result
        assert ",72," in result

    def test_negative_relative_time_clamped(self):
        """Words before clip_start should clamp to 0."""
        timings = [{"word": "early", "start": 5.0, "end": 5.5}]
        result = generate_ass_captions(timings, clip_start=10.0)
        lines = [l for l in result.strip().split("\n") if l.startswith("Dialogue:")]
        # Both start and end clamped to 0
        assert lines[0].startswith("Dialogue: 0,0:00:00.00,0:00:00.00,")


# ── File writing ─────────────────────────────────────────────────────────────

class TestWriteAssFile:
    def test_writes_content(self):
        content = "[Script Info]\ntest content\n"
        with tempfile.TemporaryDirectory() as td:
            out = write_ass_file(content, Path(td) / "sub.ass")
            assert out.exists()
            assert out.read_text() == content

    def test_creates_parent_dirs(self):
        content = "test"
        with tempfile.TemporaryDirectory() as td:
            out = write_ass_file(content, Path(td) / "nested" / "deep" / "sub.ass")
            assert out.exists()

    def test_returns_path(self):
        content = "test"
        with tempfile.TemporaryDirectory() as td:
            target = Path(td) / "sub.ass"
            result = write_ass_file(content, target)
            assert result == target