stack/tests/llm/test_pages.py

"""llm.pages — chunk → PDF page location via PyMuPDF."""

import fitz
import pytest

from llm.chunk import Chunk, Doc
from llm.pages import enrich_pdf_pages, locate, pdf_pages


@pytest.fixture
def pdf(tmp_path):
    path = tmp_path / "attachment_1.pdf"
    doc = fitz.open()
    for i, body in enumerate(
        ["Page one talks about telehealth originating sites.", "Page two covers E/M."]
    ):
        page = doc.new_page()
        page.insert_text((72, 72), f"Header {i + 1}\n{body}")
    doc.save(path)
    doc.close()
    return path


class TestPdfPages:
    def test_one_normalized_string_per_page(self, pdf):
        pages = pdf_pages(pdf)
        assert len(pages) == 2
        assert "telehealth originating sites" in pages[0]
        assert "\n" not in pages[0]

    def test_unreadable_file_is_empty(self, tmp_path):
        bad = tmp_path / "x.pdf"
        bad.write_bytes(b"not a pdf")
        assert pdf_pages(bad) == []


class TestLocate:
    def test_finds_page_by_probe(self):
        assert locate(["alpha beta gamma", "delta epsilon"], "delta epsilon") == 2

    def test_probe_normalized_before_search(self):
        assert locate(["alpha   beta\ngamma"], "alpha beta gamma") == 1

    def test_not_found_is_zero(self):
        assert locate(["alpha"], "zeta") == 0

    def test_short_probe_is_zero(self):
        assert locate(["ab cd"], "ab") == 0


class TestEnrich:
    def _chunk(self, text, section):
        return Chunk(id="c", text=text, metadata={"section": section, "seq": "0"})

    def test_sets_attachment_and_page_for_pdf_sections(self, pdf):
        doc = Doc(
            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
        )
        chunks = [
            self._chunk("Page two covers E/M.", "attachment_1.pdf"),
            self._chunk("Inline abstract text", ""),
        ]
        out = enrich_pdf_pages(doc, chunks)
        assert out[0].metadata["attachment"] == "attachment_1.pdf"
        assert out[0].metadata["page"] == "2"
        assert "attachment" not in out[1].metadata
        assert out[0].id == "c" and out[0].text == chunks[0].text

    def test_leading_section_heading_ignored_when_probing(self, pdf):
        doc = Doc(
            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
        )
        c = self._chunk(
            "## attachment_1.pdf\n\nPage one talks about telehealth originating sites.",
            "attachment_1.pdf",
        )
        (out,) = enrich_pdf_pages(doc, [c])
        assert out.metadata["page"] == "1"

    def test_unlocated_chunk_keeps_attachment_without_page(self, pdf):
        doc = Doc(
            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
        )
        (out,) = enrich_pdf_pages(
            doc, [self._chunk("nothing matches here", "attachment_1.pdf")]
        )
        assert out.metadata["attachment"] == "attachment_1.pdf"
        assert out.metadata["page"] == ""

    def test_non_pdf_section_gets_attachment_only(self, tmp_path):
        docx = tmp_path / "attachment_1.docx"
        docx.write_bytes(b"x")
        doc = Doc(
            key="K", text="", metadata={}, files=(("attachment_1.docx", str(docx)),)
        )
        (out,) = enrich_pdf_pages(doc, [self._chunk("body", "attachment_1.docx")])
        assert out.metadata["attachment"] == "attachment_1.docx"
        assert "page" not in out.metadata

    def test_no_files_is_identity(self):
        doc = Doc(key="K", text="", metadata={})
        chunks = [self._chunk("body", "attachment_1.pdf")]
        assert enrich_pdf_pages(doc, chunks) == chunks

    def test_does_not_override_existing_page(self, pdf):
        doc = Doc(
            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
        )
        c = Chunk(
            id="c",
            text="Page two covers E/M.",
            metadata={"section": "attachment_1.pdf", "page": "9"},
        )
        (out,) = enrich_pdf_pages(doc, [c])
        assert out.metadata["page"] == "9"