Spaces:
Sleeping
Sleeping
| """Tests for all enrichers and the enricher pipeline.""" | |
| from __future__ import annotations | |
| from src.app.domain.models import ( | |
| CanonicalDocument, | |
| EvidenceType, | |
| Geometry, | |
| GeometryStatus, | |
| Hyphenation, | |
| Page, | |
| Provenance, | |
| Source, | |
| TextLine, | |
| TextRegion, | |
| Word, | |
| ) | |
| from src.app.domain.models.status import InputType | |
| from src.app.enrichers import EnricherPipeline | |
| from src.app.enrichers.bbox_repair_light import BboxRepairLightEnricher | |
| from src.app.enrichers.hyphenation_basic import HyphenationBasicEnricher | |
| from src.app.enrichers.lang_propagation import LangPropagationEnricher | |
| from src.app.enrichers.polygon_to_bbox import PolygonToBboxEnricher | |
| from src.app.enrichers.reading_order_simple import ReadingOrderSimpleEnricher | |
| from src.app.enrichers.text_consistency import TextConsistencyEnricher | |
| from src.app.policies.document_policy import DocumentPolicy, strict_policy | |
| # -- Helpers ------------------------------------------------------------------ | |
| def _prov() -> Provenance: | |
| return Provenance( | |
| provider="test", adapter="v1", source_ref="$", | |
| evidence_type=EvidenceType.PROVIDER_NATIVE, | |
| ) | |
| def _geo(x: float = 10, y: float = 10, w: float = 100, h: float = 30, | |
| status: GeometryStatus = GeometryStatus.EXACT, | |
| polygon: list[tuple[float, float]] | None = None) -> Geometry: | |
| return Geometry(bbox=(x, y, w, h), polygon=polygon, status=status) | |
| def _word(wid: str, text: str, x: float = 10, y: float = 10, w: float = 50, h: float = 25, | |
| lang: str | None = None, confidence: float | None = None, | |
| status: GeometryStatus = GeometryStatus.EXACT, | |
| polygon: list[tuple[float, float]] | None = None, | |
| hyphenation: Hyphenation | None = None) -> Word: | |
| return Word( | |
| id=wid, text=text, geometry=_geo(x, y, w, h, status, polygon), | |
| provenance=_prov(), lang=lang, confidence=confidence, | |
| hyphenation=hyphenation, | |
| ) | |
| def _line(lid: str, words: list[Word], x: float = 10, y: float = 10, | |
| w: float = 200, h: float = 30, lang: str | None = None) -> TextLine: | |
| return TextLine( | |
| id=lid, geometry=_geo(x, y, w, h), provenance=_prov(), | |
| words=words, lang=lang, | |
| ) | |
| def _region(rid: str, lines: list[TextLine], x: float = 0, y: float = 0, | |
| w: float = 500, h: float = 200, lang: str | None = None) -> TextRegion: | |
| return TextRegion( | |
| id=rid, geometry=_geo(x, y, w, h), provenance=_prov(), | |
| lines=lines, lang=lang, | |
| ) | |
| def _doc(regions: list[TextRegion], width: float = 1000, height: float = 1000, | |
| reading_order: list[str] | None = None) -> CanonicalDocument: | |
| return CanonicalDocument( | |
| document_id="test", source=Source(input_type=InputType.IMAGE), | |
| pages=[Page( | |
| id="p1", page_index=0, width=width, height=height, | |
| text_regions=regions, | |
| reading_order=reading_order or [], | |
| )], | |
| ) | |
| # -- PolygonToBbox ----------------------------------------------------------- | |
| class TestPolygonToBbox: | |
| def test_derives_bbox_from_polygon(self) -> None: | |
| word = _word("w1", "test", status=GeometryStatus.UNKNOWN, | |
| polygon=[(100, 200), (400, 200), (400, 250), (100, 250)]) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])])]) | |
| enriched = PolygonToBboxEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.geometry.status == GeometryStatus.INFERRED | |
| assert w.geometry.bbox == (100, 200, 300, 50) | |
| def test_skips_when_already_exact(self) -> None: | |
| word = _word("w1", "test", status=GeometryStatus.EXACT, | |
| polygon=[(0, 0), (10, 0), (10, 10), (0, 10)]) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])])]) | |
| enriched = PolygonToBboxEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.geometry.status == GeometryStatus.EXACT | |
| def test_skips_when_no_polygon(self) -> None: | |
| word = _word("w1", "test", status=GeometryStatus.UNKNOWN) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])])]) | |
| enriched = PolygonToBboxEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.geometry.status == GeometryStatus.UNKNOWN | |
| def test_respects_policy(self) -> None: | |
| word = _word("w1", "test", status=GeometryStatus.UNKNOWN, | |
| polygon=[(0, 0), (10, 0), (10, 10), (0, 10)]) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])])]) | |
| policy = DocumentPolicy(allow_polygon_to_bbox=False) | |
| enriched = PolygonToBboxEnricher().enrich(doc, policy) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.geometry.status == GeometryStatus.UNKNOWN | |
| # -- BboxRepairLight --------------------------------------------------------- | |
| class TestBboxRepairLight: | |
| def test_clips_overflow(self) -> None: | |
| word = _word("w1", "test", x=950, y=10, w=100, h=30) # overflows 1000px page | |
| doc = _doc([_region("tb1", [_line("tl1", [word], x=900, w=200)], | |
| x=900, w=200)]) | |
| enriched = BboxRepairLightEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| # Should be clipped to page boundary | |
| assert w.geometry.bbox[0] + w.geometry.bbox[2] <= 1000 | |
| assert w.geometry.status == GeometryStatus.REPAIRED | |
| def test_no_change_when_inside(self) -> None: | |
| word = _word("w1", "test", x=10, y=10, w=50, h=25) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])])]) | |
| enriched = BboxRepairLightEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.geometry.status == GeometryStatus.EXACT | |
| assert w.geometry.bbox == (10, 10, 50, 25) | |
| # -- LangPropagation -------------------------------------------------------- | |
| class TestLangPropagation: | |
| def test_propagates_from_region(self) -> None: | |
| word = _word("w1", "hello", lang=None) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])], lang="fra")]) | |
| enriched = LangPropagationEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang == "fra" | |
| def test_propagates_from_line(self) -> None: | |
| word = _word("w1", "hello", lang=None) | |
| doc = _doc([_region("tb1", [_line("tl1", [word], lang="eng")])]) | |
| enriched = LangPropagationEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang == "eng" | |
| def test_does_not_overwrite_existing(self) -> None: | |
| word = _word("w1", "hello", lang="deu") | |
| doc = _doc([_region("tb1", [_line("tl1", [word])], lang="fra")]) | |
| enriched = LangPropagationEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang == "deu" | |
| def test_no_change_when_no_parent_lang(self) -> None: | |
| word = _word("w1", "hello", lang=None) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])])]) | |
| enriched = LangPropagationEnricher().enrich(doc, DocumentPolicy()) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang is None | |
| def test_respects_policy(self) -> None: | |
| word = _word("w1", "hello", lang=None) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])], lang="fra")]) | |
| policy = DocumentPolicy(allow_lang_propagation=False) | |
| enriched = LangPropagationEnricher().enrich(doc, policy) | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang is None | |
| # -- ReadingOrderSimple ------------------------------------------------------ | |
| class TestReadingOrderSimple: | |
| def test_orders_top_to_bottom(self) -> None: | |
| r1 = _region("tb_bottom", [_line("tl1", [_word("w1", "b")])], y=200) | |
| r2 = _region("tb_top", [_line("tl2", [_word("w2", "a")])], y=50) | |
| doc = _doc([r1, r2], reading_order=[]) | |
| enriched = ReadingOrderSimpleEnricher().enrich(doc, DocumentPolicy()) | |
| assert enriched.pages[0].reading_order == ["tb_top", "tb_bottom"] | |
| def test_skips_when_order_exists(self) -> None: | |
| r1 = _region("tb1", [_line("tl1", [_word("w1", "a")])]) | |
| doc = _doc([r1], reading_order=["tb1"]) | |
| enriched = ReadingOrderSimpleEnricher().enrich(doc, DocumentPolicy()) | |
| assert enriched.pages[0].reading_order == ["tb1"] | |
| def test_respects_policy(self) -> None: | |
| r1 = _region("tb1", [_line("tl1", [_word("w1", "a")])], y=200) | |
| r2 = _region("tb2", [_line("tl2", [_word("w2", "b")])], y=50) | |
| doc = _doc([r1, r2], reading_order=[]) | |
| enriched = ReadingOrderSimpleEnricher().enrich(doc, strict_policy()) | |
| assert enriched.pages[0].reading_order == [] | |
| # -- HyphenationBasic ------------------------------------------------------- | |
| class TestHyphenationBasic: | |
| def test_detects_hyphenation(self) -> None: | |
| line1 = _line("tl1", [_word("w1", "patri-")]) | |
| line2 = _line("tl2", [_word("w2", "moine")], y=50) | |
| doc = _doc([_region("tb1", [line1, line2])]) | |
| enriched = HyphenationBasicEnricher().enrich(doc, DocumentPolicy()) | |
| lines = enriched.pages[0].text_regions[0].lines | |
| w1 = lines[0].words[0] | |
| w2 = lines[1].words[0] | |
| assert w1.hyphenation is not None | |
| assert w1.hyphenation.is_hyphenated is True | |
| assert w1.hyphenation.part == 1 | |
| assert w1.hyphenation.full_form == "patrimoine" | |
| assert w2.hyphenation is not None | |
| assert w2.hyphenation.part == 2 | |
| assert w2.hyphenation.full_form == "patrimoine" | |
| def test_no_hyphenation_without_dash(self) -> None: | |
| line1 = _line("tl1", [_word("w1", "bonjour")]) | |
| line2 = _line("tl2", [_word("w2", "monde")], y=50) | |
| doc = _doc([_region("tb1", [line1, line2])]) | |
| enriched = HyphenationBasicEnricher().enrich(doc, DocumentPolicy()) | |
| w1 = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w1.hyphenation is None | |
| def test_no_hyphenation_if_next_uppercase(self) -> None: | |
| line1 = _line("tl1", [_word("w1", "anti-")]) | |
| line2 = _line("tl2", [_word("w2", "Américain")], y=50) | |
| doc = _doc([_region("tb1", [line1, line2])]) | |
| enriched = HyphenationBasicEnricher().enrich(doc, DocumentPolicy()) | |
| w1 = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w1.hyphenation is None | |
| def test_skips_already_hyphenated(self) -> None: | |
| hyph = Hyphenation(is_hyphenated=True, part=1, full_form="existing") | |
| line1 = _line("tl1", [_word("w1", "exist-", hyphenation=hyph)]) | |
| line2 = _line("tl2", [_word("w2", "ing")], y=50) | |
| doc = _doc([_region("tb1", [line1, line2])]) | |
| enriched = HyphenationBasicEnricher().enrich(doc, DocumentPolicy()) | |
| w1 = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w1.hyphenation.full_form == "existing" # unchanged | |
| def test_respects_policy(self) -> None: | |
| line1 = _line("tl1", [_word("w1", "patri-")]) | |
| line2 = _line("tl2", [_word("w2", "moine")], y=50) | |
| doc = _doc([_region("tb1", [line1, line2])]) | |
| enriched = HyphenationBasicEnricher().enrich(doc, strict_policy()) | |
| w1 = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w1.hyphenation is None | |
| # -- TextConsistency --------------------------------------------------------- | |
| class TestTextConsistency: | |
| def test_no_warnings_for_normal_text(self) -> None: | |
| doc = _doc([_region("tb1", [_line("tl1", [_word("w1", "hello")])])]) | |
| enriched = TextConsistencyEnricher().enrich(doc, DocumentPolicy()) | |
| assert enriched.pages[0].warnings == [] | |
| def test_warns_on_very_long_word(self) -> None: | |
| long_text = "a" * 150 | |
| doc = _doc([_region("tb1", [_line("tl1", [_word("w1", long_text)])])]) | |
| enriched = TextConsistencyEnricher().enrich(doc, DocumentPolicy()) | |
| assert any("suspiciously long" in w for w in enriched.pages[0].warnings) | |
| # -- Pipeline ---------------------------------------------------------------- | |
| class TestEnricherPipeline: | |
| def test_runs_all_enrichers(self) -> None: | |
| word = _word("w1", "hello", lang=None) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])], lang="fra")], | |
| reading_order=[]) | |
| pipeline = EnricherPipeline([ | |
| LangPropagationEnricher(), | |
| ReadingOrderSimpleEnricher(), | |
| TextConsistencyEnricher(), | |
| ]) | |
| enriched = pipeline.run(doc) | |
| # Lang propagated | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang == "fra" | |
| # Reading order inferred | |
| assert enriched.pages[0].reading_order == ["tb1"] | |
| # Audit updated | |
| assert "lang_propagation" in enriched.audit.enrichers_applied | |
| assert "reading_order_simple" in enriched.audit.enrichers_applied | |
| assert "text_consistency" in enriched.audit.enrichers_applied | |
| def test_empty_pipeline_returns_same(self) -> None: | |
| doc = _doc([_region("tb1", [_line("tl1", [_word("w1", "test")])])]) | |
| pipeline = EnricherPipeline([]) | |
| enriched = pipeline.run(doc) | |
| assert enriched.document_id == doc.document_id | |
| def test_respects_policy(self) -> None: | |
| word = _word("w1", "hello", lang=None) | |
| doc = _doc([_region("tb1", [_line("tl1", [word])], lang="fra")], | |
| reading_order=[]) | |
| pipeline = EnricherPipeline([ | |
| LangPropagationEnricher(), | |
| ReadingOrderSimpleEnricher(), | |
| ]) | |
| enriched = pipeline.run(doc, strict_policy()) | |
| # Strict policy blocks both | |
| w = enriched.pages[0].text_regions[0].lines[0].words[0] | |
| assert w.lang is None | |
| assert enriched.pages[0].reading_order == [] | |