UmbraAudio/tests/test_segmenter.py

from __future__ import annotations

from umbraaudio.mapping.segmenter import Segmenter


def test_multisentence_quote_stays_one_utterance() -> None:
    seg = Segmenter()
    text = (
        "“Fifty light-years. An old freighter like this could only jump "
        "twenty-five light-years. How can Deal Hunter do that?”"
    )
    utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
    dialogues = [u for u in utts if u["type"] == "dialogue"]
    assert len(dialogues) == 1
    assert dialogues[0]["text"].endswith("?”")
    assert "How can Deal Hunter do that?" in dialogues[0]["text"]


def test_quote_with_inner_single_quotes_not_mangled() -> None:
    seg = Segmenter()
    text = (
        "“Our best analysis would indicate that it is a message. ‘Don’t help "
        "Princess Kainda.’ Princess, we must consider how to deal with you now.”"
    )
    utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
    dialogues = [u for u in utts if u["type"] == "dialogue"]
    assert len(dialogues) == 1
    assert "‘Don’t help Princess Kainda.’" in dialogues[0]["text"]
    assert "Princess, we must consider" in dialogues[0]["text"]


def test_dialogue_and_narration_split_kept_separate() -> None:
    seg = Segmenter()
    text = "Nanny removed her helmet. “How are you feeling, Princess?”"
    utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
    types = [u["type"] for u in utts]
    assert "narration" in types
    assert "dialogue" in types


def test_carryover_speaker_after_tag() -> None:
    seg = Segmenter()
    text = "Nanny removed her helmet. “How are you feeling, Princess?”\n“I have a headache.”"
    utts = seg.segment(text, 1, character_names=["Nanny", "Kainda", "Princess"])
    dialogues = [u for u in utts if u["type"] == "dialogue"]
    assert dialogues[0]["character_name"] == "Nanny"
    assert dialogues[1]["character_name"] == "Princess"


def test_outer_quote_closed_when_inner_quote_at_end() -> None:
    seg = Segmenter()
    text = "“I am Support Robot 74532, but I have been named ‘Nanny.’”"
    utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
    dialogues = [u for u in utts if u["type"] == "dialogue"]
    assert len(dialogues) == 1
    assert dialogues[0]["text"].endswith("’”")
    assert "‘Nanny.’" in dialogues[0]["text"]


def test_inner_quote_kept_when_quote_continues() -> None:
    seg = Segmenter()
    text = "“She asked me to say, ‘Nothing will happen to you on my ship.’”"
    utts = seg.segment(text, 1, character_names=["Kainda"])
    dialogues = [u for u in utts if u["type"] == "dialogue"]
    assert len(dialogues) == 1
    assert dialogues[0]["text"] == text


def test_author_note_splits_into_sentences() -> None:
    seg = Segmenter()
    text = (
        "Author’s Note\n\n"
        "This novel, Deal Hunter, does not use the terms android or droid. "
        "Google uses “Android” to refer to the operating system on phones.\n\n"
        "Deal Hunter uses the generic terms of robot and artificial intelligence (AI)."
    )
    utts = seg.segment(text, 5, chapter_title="Author's Note")
    texts = [u["text"] for u in utts]
    assert len(texts) == 4
    assert texts[0] == "Author’s Note"
    assert "does not use the terms android or droid." in texts[1]
    assert "Google uses “Android” to refer to the operating system on phones." in texts[2]
    assert texts[3].startswith("Deal Hunter uses the generic terms")


def test_author_note_all_narration() -> None:
    seg = Segmenter()
    text = (
        "Author’s Note\n\n"
        "Every star has humans, and they are looking for resources. "
        "The Dial sector has several pirate clans."
    )
    utts = seg.segment(text, 5, chapter_title="Author's Note")
    assert all(u["type"] == "narration" for u in utts)
    assert all(u["character_name"] == "narrator" for u in utts)


def test_author_note_merges_abbreviations() -> None:
    seg = Segmenter()
    text = "Author’s Note\n\nThis is a work of fiction. Dr. Smith wrote it."
    utts = seg.segment(text, 5, chapter_title="Author's Note")
    texts = [u["text"] for u in utts]
    assert "This is a work of fiction." in texts
    assert "Dr. Smith wrote it." in texts


def test_unclosed_quote_does_not_swallow_following_narration() -> None:
    seg = Segmenter()
    text = (
        "As she says this, there is a roar, and Niñera says, “The captain is "
        "changing course, and that was a missile he is firing at the pirates. \n"
        "As Niñera helps Kainda, they feel the impacts more, and the air "
        "pressure alarm sounds. The computer voice announces, “Hull breach. "
        "Please put on emergency space suits.”\nSealing the suit, the robot "
        "picks up the helmet."
    )
    utts = seg.segment(text, 6, character_names=["Niñera", "Kainda"])
    assert len(utts) >= 3
    joined = " ".join(u["text"] for u in utts)
    assert "Hull breach" in joined
    hull = [u for u in utts if "Hull breach" in u["text"]]
    assert len(hull) == 1
    assert "changing course" not in hull[0]["text"]


def test_unterminated_quote_at_eof_becomes_plain_text() -> None:
    seg = Segmenter()
    text = "She stared out the viewport. “We are never going to make it"
    utts = seg.segment(text, 1)
    assert all(u["type"] == "narration" for u in utts)
    assert any("never going to make it" in u["text"] for u in utts)


def test_balanced_multisentence_quote_still_intact() -> None:
    seg = Segmenter()
    text = "He said, “Line one. Line two. Line three.” Then he left."
    utts = seg.segment(text, 1)
    dialogues = [u for u in utts if u["type"] == "dialogue"]
    assert len(dialogues) == 1
    assert dialogues[0]["text"].startswith("“")
    assert dialogues[0]["text"].endswith("”")