abogen/tests/test_subtitle_generation.py

"""Tests for abogen.domain.subtitle_generation module."""

import pytest

from abogen.domain.subtitle_generation import (
    process_subtitle_tokens,
    PUNCTUATION_SENTENCE,
    PUNCTUATION_SENTENCE_COMMA,
)


class TestProcessSubtitleTokens:
    """Tests for process_subtitle_tokens function."""

    def test_empty_tokens(self):
        """Test processing empty token list does nothing."""
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=[],
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Sentence",
            lang_code="a",
        )
        assert entries == []

    def test_disabled_mode(self):
        """Test Disabled mode returns no entries."""
        tokens = [
            {"start": 0.0, "end": 1.0, "text": "Hello", "whitespace": " "},
            {"start": 1.0, "end": 2.0, "text": "world", "whitespace": ""},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Disabled",
            lang_code="a",
        )
        assert entries == []

    def test_line_mode_basic(self):
        """Test Line mode splits on newlines."""
        tokens = [
            {"start": 0.0, "end": 1.0, "text": "First line", "whitespace": "\n"},
            {"start": 1.0, "end": 2.0, "text": "Second line", "whitespace": ""},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Line",
            lang_code="a",
        )
        assert len(entries) == 2
        assert entries[0][2] == "First line"
        assert entries[1][2] == "Second line"

    def test_sentence_mode_punctuation_split(self):
        """Test Sentence mode splits on sentence punctuation."""
        tokens = [
            {"start": 0.0, "end": 0.5, "text": "First sentence", "whitespace": " "},
            {"start": 0.5, "end": 1.0, "text": ".", "whitespace": " "},
            {"start": 1.0, "end": 1.5, "text": "Second sentence", "whitespace": " "},
            {"start": 1.5, "end": 2.0, "text": ".", "whitespace": ""},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Sentence",
            lang_code="a",
        )
        assert len(entries) >= 1
        # Should have at least one entry with both sentences or split
        combined_text = " ".join(e[2] for e in entries)
        assert "First sentence" in combined_text
        assert "Second sentence" in combined_text

    def test_word_count_mode(self):
        """Test word count mode (e.g., '5' for 5 words per entry)."""
        tokens = [
            {"start": 0.0, "end": 0.2, "text": "word1", "whitespace": " "},
            {"start": 0.2, "end": 0.4, "text": "word2", "whitespace": " "},
            {"start": 0.4, "end": 0.6, "text": "word3", "whitespace": " "},
            {"start": 0.6, "end": 0.8, "text": "word4", "whitespace": " "},
            {"start": 0.8, "end": 1.0, "text": "word5", "whitespace": " "},
            {"start": 1.0, "end": 1.2, "text": "word6", "whitespace": " "},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="2",  # 2 words per entry
            lang_code="a",
        )
        assert len(entries) >= 2
        # Check that entries are split roughly by word count
        for entry in entries:
            # Each entry should have at least one word
            assert len(entry[2].split()) >= 1

    def test_fallback_end_time(self):
        """Test fallback_end_time is applied when end time is invalid."""
        tokens = [
            {"start": 0.0, "end": None, "text": "Test", "whitespace": ""},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Line",
            lang_code="a",
            fallback_end_time=10.0,
        )
        assert len(entries) == 1
        assert entries[0][1] == 10.0  # Should use fallback

    def test_karaoke_highlighting_mode(self):
        """Test Sentence + Highlighting mode generates karaoke tags."""
        tokens = [
            {"start": 0.0, "end": 0.5, "text": "Hello", "whitespace": " "},
            {"start": 0.5, "end": 1.0, "text": "world", "whitespace": ""},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Sentence + Highlighting",
            lang_code="a",
        )
        assert len(entries) >= 1
        # Should contain karaoke tags
        text = entries[0][2]
        assert "{\\kf" in text

    def test_max_subtitle_words_limit(self):
        """Test that max_subtitle_words limits entry length."""
        tokens = [
            {"start": float(i), "end": float(i + 0.1), "text": f"word{i}", "whitespace": " "}
            for i in range(10)
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=3,
            subtitle_mode="Line",
            lang_code="a",
        )
        # Should have more than 1 entry due to word limit
        assert len(entries) > 1

    def test_preserves_token_timing(self):
        """Test that token timing is preserved in entries."""
        tokens = [
            {"start": 0.0, "end": 1.0, "text": "First", "whitespace": " "},
            {"start": 1.0, "end": 2.0, "text": "Second", "whitespace": ""},
        ]
        entries = []
        process_subtitle_tokens(
            tokens_with_timestamps=tokens,
            subtitle_entries=entries,
            max_subtitle_words=50,
            subtitle_mode="Sentence",
            lang_code="a",
        )
        assert len(entries) >= 1
        # Check that timing is preserved
        for entry in entries:
            assert entry[0] >= 0.0
            assert entry[1] >= entry[0]


class TestPunctuationConstants:
    """Tests for punctuation constants."""

    def test_punctuation_sentence_contains_basic(self):
        """Test PUNCTUATION_SENTENCE contains basic sentence punctuation."""
        assert "." in PUNCTUATION_SENTENCE
        assert "!" in PUNCTUATION_SENTENCE
        assert "?" in PUNCTUATION_SENTENCE

    def test_punctuation_sentence_comma_contains_comma(self):
        """Test PUNCTUATION_SENTENCE_COMMA contains comma."""
        assert "," in PUNCTUATION_SENTENCE_COMMA
        assert "." in PUNCTUATION_SENTENCE_COMMA
        assert "!" in PUNCTUATION_SENTENCE_COMMA
        assert "?" in PUNCTUATION_SENTENCE_COMMA