tests/unit/test_pdf_layout.py

"""
Unit tests for layout-aware PDF extraction.

Two capabilities are covered:
  * tables survive extraction as Markdown instead of being flattened into an
    unreadable run of cell values
  * pages with no extractable text are routed to OCR, within a budget, and
    degrade quietly when Tesseract is absent
"""

import os
from unittest.mock import patch

import pytest

from app.services.pdf_layout import PageExtraction, extract_page, table_to_markdown

pytestmark = pytest.mark.unit


# ── Table rendering ────────────────────────────────────────────────────────────


def test_table_renders_as_markdown_with_header_separator():
    md = table_to_markdown([["Height", "Bolt"], ["4 ft", "36 in"], ["6 ft", "48 in"]])
    lines = md.splitlines()
    assert lines[0] == "| Height | Bolt |"
    assert lines[1] == "| --- | --- |"
    assert lines[2] == "| 4 ft | 36 in |"


def test_row_column_relationship_is_preserved():
    """The whole point: a value must stay associated with its row and column."""
    md = table_to_markdown(
        [
            ["Mining Height", "Bolt Length"],
            ["under 4 ft", "36 in"],
            ["6 to 8 ft", "60 in"],
        ]
    )
    # The 60 in row must be readable as belonging to 6 to 8 ft.
    assert "| 6 to 8 ft | 60 in |" in md


def test_none_cells_become_empty_not_the_string_none():
    md = table_to_markdown([["A", "B"], ["x", None]])
    assert "None" not in md
    assert "| x |  |" in md


def test_ragged_rows_are_padded():
    md = table_to_markdown([["A", "B", "C"], ["1"], ["1", "2"]])
    for line in md.splitlines():
        assert line.count("|") == 4


def test_embedded_pipes_cannot_break_the_table():
    md = table_to_markdown([["Header"], ["a | b"]])
    assert r"a \| b" in md


def test_embedded_newlines_are_flattened():
    md = table_to_markdown([["Header"], ["line1\nline2"]])
    assert "line1 line2" in md
    assert len(md.splitlines()) == 3


def test_fully_empty_rows_are_dropped():
    md = table_to_markdown([["A", "B"], ["", ""], ["1", "2"]])
    assert len(md.splitlines()) == 3  # header, separator, one data row


def test_empty_table_returns_empty_string():
    assert table_to_markdown([]) == ""
    assert table_to_markdown([["", ""]]) == ""


# ── Page extraction / OCR routing ──────────────────────────────────────────────


class FakePage:
    def __init__(self, text="", tables=None):
        self._text = text
        self._tables = tables or []

    def extract_text(self):
        return self._text

    def extract_tables(self):
        return self._tables


def test_text_page_is_not_sent_to_ocr():
    page = FakePage(text="x" * 500)
    with patch("app.services.pdf_layout.ocr_page") as ocr:
        result = extract_page(page, "f.pdf", 1, ocr_budget=[10])
    ocr.assert_not_called()
    assert not result.used_ocr


def test_scanned_page_is_sent_to_ocr():
    """A page yielding nothing is the signature of a scan."""
    page = FakePage(text="")
    with patch("app.services.pdf_layout.ocr_page", return_value="RECOVERED TEXT"):
        result = extract_page(page, "f.pdf", 3, ocr_budget=[10])
    assert result.used_ocr
    assert "RECOVERED TEXT" in result.text


def test_ocr_budget_is_consumed_and_enforced():
    page = FakePage(text="")
    budget = [2]
    with patch("app.services.pdf_layout.ocr_page", return_value="TEXT"):
        assert extract_page(page, "f.pdf", 1, budget).used_ocr
        assert extract_page(page, "f.pdf", 2, budget).used_ocr
        # Budget exhausted — third scanned page is not OCR'd.
        assert not extract_page(page, "f.pdf", 3, budget).used_ocr
    assert budget == [0]


def test_ocr_failure_degrades_quietly():
    """No Tesseract binary must not raise — the document still ingests."""
    page = FakePage(text="")
    with patch("app.services.pdf_layout.ocr_page", return_value=""):
        result = extract_page(page, "f.pdf", 1, ocr_budget=[5])
    assert result.used_ocr is False
    assert result.text == ""


def test_ocr_supplements_rather_than_replaces_sparse_text():
    page = FakePage(text="Figure 4.")
    with patch("app.services.pdf_layout.ocr_page", return_value="scanned body"):
        result = extract_page(page, "f.pdf", 1, ocr_budget=[5])
    assert "Figure 4." in result.text
    assert "scanned body" in result.text


def test_tables_are_appended_and_labelled_with_page():
    page = FakePage(text="Prose here.", tables=[[["A", "B"], ["1", "2"]]])
    result = extract_page(page, "f.pdf", 7, ocr_budget=[5])
    assert "Prose here." in result.text
    assert "[Table 1 on page 7]" in result.text
    assert "| A | B |" in result.text
    assert result.tables_found == 1


def test_table_extraction_error_does_not_fail_the_page():
    class Exploding(FakePage):
        def extract_tables(self):
            raise RuntimeError("malformed table object")

    result = extract_page(Exploding(text="Prose"), "f.pdf", 1, ocr_budget=[5])
    assert result.text == "Prose"
    assert result.tables_found == 0


def test_table_can_rescue_a_page_from_being_treated_as_scanned():
    """A page that is only a table has little prose but is not a scan."""
    big_table = [["Col"] * 4] + [[f"v{i}{j}" for j in range(4)] for i in range(12)]
    page = FakePage(text="", tables=[big_table])
    with patch("app.services.pdf_layout.ocr_page") as ocr:
        result = extract_page(page, "f.pdf", 1, ocr_budget=[5])
    ocr.assert_not_called()
    assert result.tables_found == 1


# ── End-to-end against a real PDF ──────────────────────────────────────────────

FIXTURE = os.path.join(
    os.path.dirname(__file__), "..", "fixtures", "roof_control_table.pdf"
)


@pytest.mark.skipif(not os.path.exists(FIXTURE), reason="fixture PDF not built")
def test_real_pdf_table_survives_extraction():
    pytest.importorskip("pdfplumber")
    from app.services.extractors import PDFExtractor

    doc = PDFExtractor().extract(FIXTURE)

    assert "| Mining Height | Bolt Length | Spacing | Support Type |" in doc.full_text
    assert "| 6 to 8 ft | 60 in | 3.5 ft | Mechanical |" in doc.full_text
    assert doc.metadata["tables_extracted"] == 1


@pytest.mark.skipif(not os.path.exists(FIXTURE), reason="fixture PDF not built")
def test_pypdf_fallback_loses_the_table_structure():
    """
    Documents the regression this replaces: the old path cannot answer
    "what bolt length for 6 to 8 ft?" because the association is gone.
    """
    from app.services.extractors import PDFExtractor

    text = PDFExtractor()._extract_with_pypdf(FIXTURE).full_text
    assert "| 6 to 8 ft | 60 in |" not in text