path: root/app/tests/test_chunking.py

"""Tests for text chunking."""

import unittest
from unittest.mock import patch

from converter import config
from converter.chunking import split_into_chunks


class ChunkSizeDefaultTests(unittest.TestCase):
    """Guard the request-size setting: each API call is one model
    generation, and the servers silently truncate audio when a single
    generation runs too long (~2.5 min faster backend, ~11 min Qwen
    demo), so the default chunk size must stay well inside that budget.
    There is no hard ceiling beyond CHUNK_SIZE; users raising it accept
    the truncation risk themselves."""

    def test_default_chunk_size_within_single_generation_budget(self):
        self.assertLessEqual(config.CHUNK_SIZE, 300)

    def test_default_chunk_size_is_positive(self):
        self.assertGreaterEqual(config.CHUNK_SIZE, 1)


class RequestSizeTests(unittest.TestCase):
    def test_oversized_chunk_size_is_honored(self):
        # No clamping: whatever size is configured (or requested) is used.
        text = " ".join(f"word{i}" for i in range(30)) + "."
        chunks = split_into_chunks(text, max_words=5000)
        self.assertEqual(len(chunks), 1)
        self.assertEqual(len(chunks[0].split()), 30)

    def test_default_uses_runtime_config_chunk_size(self):
        # The default resolves config.CHUNK_SIZE at call time, so
        # patching the config changes the default split size.
        sentences = " ".join(
            f"S{i} " + " ".join(["word"] * 8) + "." for i in range(60))
        with patch.object(config, "CHUNK_SIZE", 120):
            chunks = split_into_chunks(sentences)
        self.assertGreater(len(chunks), 1)
        self.assertTrue(all(len(chunk.split()) <= 120 for chunk in chunks))


class SplitIntoChunksTests(unittest.TestCase):
    def test_empty_input(self):
        self.assertEqual(split_into_chunks(""), [])
        self.assertEqual(split_into_chunks("   \n  "), [])

    def test_short_text_single_chunk(self):
        self.assertEqual(split_into_chunks("One short sentence."), ["One short sentence."])

    def test_respects_word_limit_across_sentences(self):
        # 10 sentences of 9 words each = 90 words total
        sentences = [f"S{i} " + " ".join(["word"] * 8) + "." for i in range(10)]
        chunks = split_into_chunks(" ".join(sentences), max_words=25)
        self.assertGreater(len(chunks), 1)
        self.assertTrue(all(len(c.split()) <= 25 for c in chunks))
        self.assertEqual(sum(len(c.split()) for c in chunks), 90)

    def test_long_sentence_split_keeps_punctuation(self):
        # 10 clauses of 5 words each, joined by comma+space
        sentence = ", ".join([" ".join(["w"] * 5) for _ in range(10)]) + "."
        chunks = split_into_chunks(sentence, max_words=12)
        self.assertGreater(len(chunks), 1)
        self.assertTrue(all(len(c.split()) <= 12 for c in chunks))
        self.assertIn(",", chunks[0])  # commas retained for TTS prosody

    def test_clause_split_never_breaks_numbers(self):
        # Regression: the clause split used to fire at every comma even
        # without whitespace, mutating "1,000,000" into "1, 000, 000".
        sentence = ("There were exactly 1,000,000 soldiers marching at 12:30, "
                    + "and they kept marching onward " * 30) + "endlessly."
        chunks = split_into_chunks(sentence, max_words=25)
        self.assertGreater(len(chunks), 1)
        joined = " ".join(chunks)
        self.assertIn("1,000,000", joined)
        self.assertIn("12:30", joined)
        self.assertNotIn("1, 000", joined)
        self.assertNotIn("000, 000", joined)
        self.assertNotIn("12: 30", joined)

    def test_clause_split_requires_whitespace_after_punctuation(self):
        # Run-on clauses without spaces after commas have no clause split
        # point, so the last-resort word-boundary split fires instead.
        # Tokens themselves (and numbers like "1,000,000") stay intact.
        sentence = ",".join([" ".join(["w"] * 5) for _ in range(10)]) + "."
        chunks = split_into_chunks(sentence, max_words=12)
        self.assertGreater(len(chunks), 1)
        self.assertTrue(all(len(chunk.split()) <= 12 for chunk in chunks))
        tokens = sentence.replace(",", " , ").split()
        rejoined = " ".join(chunks).replace(",", " , ").split()
        self.assertEqual(rejoined, tokens)

    def test_single_oversized_sentence_is_word_split(self):
        # A punctuation-free sentence longer than the limit is split at word
        # boundaries so no single request exceeds the configured size.
        sentence = " ".join(["word"] * 30) + "."
        chunks = split_into_chunks(sentence, max_words=10)
        self.assertGreater(len(chunks), 1)
        self.assertTrue(all(len(chunk.split()) <= 10 for chunk in chunks))
        self.assertEqual(sum(len(chunk.split()) for chunk in chunks), 30)

    def test_word_split_never_breaks_number_tokens(self):
        # Numbers and other punctuation-bearing tokens are single words and
        # must never be broken apart by the last-resort word split.
        sentence = ("There were exactly 1,000,000 soldiers marching at 12:30 "
                    "and " + "they kept marching onward " * 20) + "endlessly."
        chunks = split_into_chunks(sentence, max_words=10)
        self.assertGreater(len(chunks), 1)
        joined = " ".join(chunks)
        self.assertIn("1,000,000", joined)
        self.assertIn("12:30", joined)
        self.assertNotIn("1, 000", joined)
        self.assertNotIn("12: 30", joined)

    def test_smart_strategy_distributes_sentences(self):
        # The smart strategy leaves whole sentences intact: a chunk
        # crossing the target (85% of the limit) ends at the next clean
        # boundary early; legacy keeps packing until the limit.
        sentences = " ".join(
            " ".join(f"Sw{j}" for j in range(n)) + "."
            for n in (9, 8, 3, 9, 8, 3))
        smart = split_into_chunks(sentences, max_words=20, smart=True)
        legacy = split_into_chunks(sentences, max_words=20, smart=False)
        self.assertEqual([len(c.split()) for c in smart], [17, 20, 3])
        self.assertEqual([len(c.split()) for c in legacy], [20, 20])
        self.assertEqual(self._content(smart),
                         self._content(sentences.split(" ")))

    def test_dialogue_tag_stays_with_its_quote(self):
        # A closing quote followed by a lowercase attribution ("she
        # said.") is one sentence unit: the chunk never separates them.
        text = "“Come here!” she said. “Are you coming?” He did not answer."
        chunks = split_into_chunks(text, max_words=10, smart=True)
        self.assertIn("“Come here!” she said.", chunks[0])
        self._assert_bounded(chunks, 10)
        self._assert_content(text, chunks)

    def test_abbreviations_and_initials_do_not_split(self):
        text = ("Dr. Smith met J. K. R. at the U.S. border with No. 3 "
                "shortly. He continued onward.")
        chunks = split_into_chunks(text, max_words=50, smart=True)
        self.assertEqual(len(chunks), 1)
        self.assertIn("Dr. Smith", chunks[0])
        self.assertIn("J. K. R.", chunks[0])
        self.assertIn("U.S.", chunks[0])
        self._assert_content(text, chunks)

    def test_decimals_are_never_boundaries(self):
        text = "Pi is 3.14159 and the toll was 1,000,000 miles. Next one."
        chunks = split_into_chunks(text, max_words=12, smart=True)
        joined = " ".join(chunks)
        self.assertIn("3.14159", joined)
        self.assertIn("1,000,000", joined)
        self._assert_bounded(chunks, 12)
        self._assert_content(text, chunks)

    def test_unclosed_quote_still_respects_the_limit(self):
        # A quote with no closing mark must not grow chunks or silence
        # later boundaries; units carry the open state and stay bounded.
        text = "“No closing quote follows. " + "filler words here. " * 20
        chunks = split_into_chunks(text, max_words=10, smart=True)
        self._assert_bounded(chunks, 10)
        self._assert_content(text, chunks)

    def test_paragraph_boundary_resets_an_open_quote(self):
        # Blank lines reset the quotation state: an unclosed quote in
        # one paragraph cannot trap later paragraphs inside the quote.
        text = ("“Still open. filler words inside the quote here.\n\n"
                "New paragraph. Words outside any quote now, plenty.")
        chunks = split_into_chunks(text, max_words=10, smart=True)
        self._assert_bounded(chunks, 10)
        self._assert_content(text, chunks)
        second = [c for c in chunks if "New paragraph." in c]
        self.assertTrue(second)
        # The paragraph's text starts a fresh chunk: the open-quote
        # prefix did not swallow it.
        self.assertTrue(
            any(chunk.startswith("New paragraph.") or
                "“Still open." not in chunk
                for chunk in chunks))

    def test_quote_retreat_keeps_short_quotes_whole(self):
        # When the next sentence would overflow, the break retreats to
        # the last boundary that ended outside a quotation — before the
        # quote — so a short quotation lands in one request.
        text = ("Plain opening words here to fill. "
                "“A short quote. With two sentences. Inside.” "
                "More trailing prose follows this quote. And more.")
        chunks = split_into_chunks(text, max_words=14, smart=True)
        self._assert_bounded(chunks, 14)
        self._assert_content(text, chunks)
        quote_chunks = [c for c in chunks if "short quote" in c]
        self.assertTrue(quote_chunks)
        self.assertIn("Inside.", quote_chunks[0])

    def test_oversized_single_tokens_respect_char_cap(self):
        # Degenerate text (punctuation-free run): bounded by the hard
        # character cap, with every character preserved.
        text = "a" * 5000
        chunks = split_into_chunks(text, max_words=250, smart=True)
        self.assertGreater(len(chunks), 1)
        self.assertTrue(all(len(c) <= 250 * 8 for c in chunks))
        self.assertEqual("".join(chunks), text)

    def test_non_spaced_script_stays_bounded(self):
        # CJK text: sentence terminators break units without
        # whitespace, and each ideograph counts as a word.
        text = "。".join(["字" * 15 for _ in range(30)]) + "。"
        chunks = split_into_chunks(text, max_words=40, smart=True)
        self._assert_bounded(chunks, 40)
        self._assert_content(text, chunks)
        joined = " ".join(chunks).split()
        self.assertTrue(all(len(j) <= 33 for j in joined))

    def test_legacy_and_smart_matching_smart_negation(self):
        # smart=False preserves the legacy splitting exactly.
        sentences = " ".join(
            " ".join(f"Sw{j}" for j in range(9)) + "." for _ in range(10))
        with patch.object(config, "CHUNK_SIZE", 25):
            legacy_default = split_into_chunks(sentences, smart=False)
        self.assertEqual([len(c.split()) for c in legacy_default],
                         [18] * 5)
        self.assertEqual(self._content(legacy_default),
                         self._content(sentences.split()))

    # -- helpers ---------------------------------------------------------

    @staticmethod
    def _content(chunks_or_words):
        """Non-whitespace content of the pieces (or of a token list)."""
        pieces = (chunks_or_words if isinstance(chunks_or_words, str)
                  else " ".join(chunks_or_words))
        return "".join(pieces.split())

    def _assert_content(self, text, chunks):
        self.assertEqual(self._content(chunks), self._content(text))

    def _assert_bounded(self, chunks, max_words):
        from converter.chunking import _count_words
        for chunk in chunks:
            self.assertLessEqual(_count_words(chunk), max_words,
                                 chunk[:40])


if __name__ == "__main__":
    unittest.main()