ml_pipeline/tests/test_segmenter.py

"""
test_segmenter.py
=================
Unit tests for StructuralSegmenter.

Run with:
  cd ai-service && python -m pytest ml_pipeline/tests/test_segmenter.py -v
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))

import pytest
from ml_pipeline.structural_segmenter import StructuralSegmenter


# ── Helpers ─────────────────────────────────────────────────────────────────────

def _span(text: str, size: float = 12.0, flags: int = 0) -> dict:
    return {"text": text, "size": size, "flags": flags, "font": "Helvetica"}


def _block(lines_texts: list[str], size: float = 12.0, flags: int = 0, page: int = 0) -> dict:
    """Build a single PyMuPDF block dict."""
    return {
        "type": 0,
        "_page": page,
        "lines": [
            {"spans": [_span(t, size, flags)], "bbox": [0, float(i * 14), 400, float(i * 14 + 12)]}
            for i, t in enumerate(lines_texts)
        ],
    }


# ── Tests: Play segmentation ─────────────────────────────────────────────────────

class TestPlaySegmentation:
    seg = StructuralSegmenter()

    def _blocks(self):
        return [
            _block(["ACT I"],    size=20, flags=16),   # large + bold → act heading
            _block(["SCENE 1"],  size=16, flags=16),   # large + bold → scene heading
            _block(["HAMLET.",   "To be, or not to be, that is the question:"]),
            _block(["[Enter HORATIO]"]),
            _block(["HORATIO.", "O, farewell, honest soldier."]),
            _block(["SCENE 2"],  size=16, flags=16),
            _block(["OPHELIA.", "Good night, sweet prince."]),
            _block(["ACT II"],   size=20, flags=16),
            _block(["SCENE 1"],  size=16, flags=16),
            _block(["MARCELLUS.", "Something is rotten in the state of Denmark."]),
        ]

    def test_produces_two_acts(self):
        units = self.seg.segment(self._blocks(), doc_type="play")
        assert len(units) == 2
        assert units[0]["title"].upper().startswith("ACT I")
        assert units[1]["title"].upper().startswith("ACT II")

    def test_act_contains_children_scenes(self):
        units = self.seg.segment(self._blocks(), doc_type="play")
        act1 = units[0]
        assert "children" in act1
        assert len(act1["children"]) >= 1

    def test_scene_contains_blocks(self):
        units = self.seg.segment(self._blocks(), doc_type="play")
        scene1 = units[0]["children"][0]
        assert "blocks" in scene1
        assert len(scene1["blocks"]) > 0

    def test_dialogue_block_has_character(self):
        units = self.seg.segment(self._blocks(), doc_type="play")
        # Find any dialogue block
        for act in units:
            for scene in act.get("children", []):
                for block in scene.get("blocks", []):
                    if block["type"] == "dialogue":
                        assert block["character"] is not None
                        assert len(block["character"]) > 0
                        return
        pytest.fail("No dialogue block found in output")

    def test_stage_direction_type(self):
        units = self.seg.segment(self._blocks(), doc_type="play")
        for act in units:
            for scene in act.get("children", []):
                for block in scene.get("blocks", []):
                    if block["type"] == "stage_direction":
                        assert block["character"] is None
                        return

    def test_units_have_ids(self):
        units = self.seg.segment(self._blocks(), doc_type="play")
        for act in units:
            assert "id" in act and act["id"]
            for scene in act.get("children", []):
                assert "id" in scene and scene["id"]

    def test_act_without_explicit_heading_creates_inferred(self):
        # Only scene headings, no ACT heading → should still produce hierarchy
        blocks = [
            _block(["SCENE 1"],  size=16, flags=16),
            _block(["ROMEO.", "But soft, what light through yonder window breaks?"]),
        ]
        units = self.seg.segment(blocks, doc_type="play")
        assert len(units) >= 1
        assert units[0]["inferred"] is True

    def test_empty_blocks_returns_list(self):
        units = self.seg.segment([], doc_type="play")
        assert isinstance(units, list)


# ── Tests: Novel segmentation ────────────────────────────────────────────────────

class TestNovelSegmentation:
    seg = StructuralSegmenter()

    def _blocks(self):
        return [
            _block(["Chapter 1"], size=18, flags=16),
            _block(["It was the best of times, it was the worst of times, it was the age of wisdom, "
                    "it was the age of foolishness, it was the epoch of belief."]),
            _block(["It was a far, far better thing that I do, than I have ever done before."]),
            _block(["Chapter 2"], size=18, flags=16),
            _block(["The great fish moved silently through the night water, propelled by short sweeps "
                    "of its crescent tail."]),
        ]

    def test_produces_two_chapters(self):
        units = self.seg.segment(self._blocks(), doc_type="novel")
        assert len(units) == 2

    def test_chapter_has_children(self):
        units = self.seg.segment(self._blocks(), doc_type="novel")
        assert "children" in units[0]
        assert len(units[0]["children"]) >= 1

    def test_chapter_content_rolled_up(self):
        units = self.seg.segment(self._blocks(), doc_type="novel")
        # Chapter-level content should be non-empty
        assert len(units[0]["content"]) > 0

    def test_section_has_paragraphs(self):
        units = self.seg.segment(self._blocks(), doc_type="novel")
        for ch in units:
            for sec in ch.get("children", []):
                assert "paragraphs" in sec
                assert isinstance(sec["paragraphs"], list)

    def test_generic_type_produces_chapters(self):
        blocks = self._blocks()
        units = self.seg.segment(blocks, doc_type="generic")
        assert isinstance(units, list)
        assert len(units) >= 1


# ── Tests: Heading threshold ─────────────────────────────────────────────────────

class TestHeadingThreshold:
    seg = StructuralSegmenter()

    def test_threshold_above_median(self):
        from ml_pipeline.structural_segmenter import Span
        spans = [
            Span("body", 12.0, 0, 0),
            Span("body", 12.0, 0, 0),
            Span("body", 11.5, 0, 0),
            Span("heading", 20.0, 16, 0),
        ]
        threshold = self.seg._compute_heading_threshold(spans)
        assert threshold > 12.0   # Must be above body text size
        assert threshold < 30.0   # Must be below an obviously unreasonable ceiling

    def test_threshold_with_single_span(self):
        from ml_pipeline.structural_segmenter import Span
        spans = [Span("only", 12.0, 0, 0)]
        threshold = self.seg._compute_heading_threshold(spans)
        assert threshold > 0