lllin000_PaperForge/tests/test_ocr_document.py

89 lines
3.8 KiB
Python

from __future__ import annotations
def test_analyze_document_structure_flat_backmatter() -> None:
from paperforge.worker.ocr_document import DocumentStructure, PagePosition, analyze_document_structure
blocks = [
{"page": 1, "role": "body_paragraph", "text": "Intro"},
{"page": 2, "role": "body_paragraph", "text": "Methods"},
{"page": 3, "role": "body_paragraph", "text": "Results"},
{"page": 3, "role": "section_heading", "text": "Discussion"},
{"page": 4, "role": "body_paragraph", "text": "Discussion text"},
{"page": 5, "role": "backmatter_heading", "text": "Acknowledgments"},
{"page": 5, "role": "backmatter_body", "text": "Thanks."},
{"page": 6, "role": "reference_heading", "text": "References"},
{"page": 6, "role": "reference_item", "text": "[1] Ref A"},
]
doc = analyze_document_structure(blocks)
assert isinstance(doc, DocumentStructure)
assert doc.body_end_page == 4
assert doc.spread_start == 5
assert doc.spread_end == 6
assert doc.backmatter_form == "flat"
def test_analyze_document_structure_container_backmatter() -> None:
from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure
blocks = [
{"page": 1, "role": "body_paragraph", "text": "Intro"},
{"page": 2, "role": "section_heading", "text": "Methods"},
{"page": 3, "role": "body_paragraph", "text": "Results"},
{"page": 3, "role": "section_heading", "text": "Discussion"},
{"page": 4, "role": "body_paragraph", "text": "Discussion text"},
{"page": 5, "role": "backmatter_boundary_heading", "text": "ADDITIONAL INFORMATION AND DECLARATIONS"},
{"page": 5, "role": "backmatter_heading", "text": "Funding"},
{"page": 5, "role": "backmatter_body", "text": "Funded by NIH."},
{"page": 6, "role": "backmatter_heading", "text": "Acknowledgments"},
{"page": 6, "role": "backmatter_body", "text": "Thanks."},
{"page": 6, "role": "backmatter_heading", "text": "Author Contributions"},
{"page": 6, "role": "backmatter_body", "text": "AB did X."},
{"page": 7, "role": "reference_heading", "text": "References"},
{"page": 7, "role": "reference_item", "text": "[1] Ref A"},
]
doc = analyze_document_structure(blocks)
assert isinstance(doc, DocumentStructure)
assert doc.backmatter_form == "container"
def test_analyze_document_structure_no_tail() -> None:
from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure
blocks = [
{"page": 1, "role": "body_paragraph", "text": "Only body"},
{"page": 2, "role": "body_paragraph", "text": "Still body"},
]
doc = analyze_document_structure(blocks)
assert isinstance(doc, DocumentStructure)
assert doc.body_end_page is None
assert doc.spread_start is None
assert doc.spread_end is None
assert doc.backmatter_form == "flat"
def test_analyze_document_structure_mixed_tail_spread() -> None:
from paperforge.worker.ocr_document import DocumentStructure, analyze_document_structure
blocks = [
{"page": 1, "role": "body_paragraph", "text": "Intro"},
{"page": 2, "role": "body_paragraph", "text": "Results"},
{"page": 3, "role": "body_paragraph", "text": "Continued results"},
{"page": 3, "role": "backmatter_heading", "text": "Acknowledgments"},
{"page": 3, "role": "backmatter_body", "text": "Thanks."},
{"page": 4, "role": "body_paragraph", "text": "More results interleaved"},
{"page": 4, "role": "reference_heading", "text": "References"},
{"page": 4, "role": "reference_item", "text": "[1] Ref A"},
]
doc = analyze_document_structure(blocks)
assert isinstance(doc, DocumentStructure)
assert doc.spread_start is not None
assert doc.spread_end is not None