from __future__ import annotations import hashlib import io import json from contextlib import contextmanager from datetime import UTC, datetime from pathlib import Path import httpx from PIL import Image from reportlab.lib.pdfencrypt import StandardEncryption from reportlab.lib.utils import ImageReader from reportlab.pdfgen import canvas from gcmd_classifier.config import DatasetExtractionSettings from gcmd_classifier.datasets.cmr import resolve_cmr_collection, validate_cmr_url from gcmd_classifier.datasets.documents import RetrievedREADME, retrieve_selected_readme from gcmd_classifier.datasets.extraction import extract_pdf_manifest from gcmd_classifier.datasets.models import ( BlindCollectionView, CMRRetrievedSource, DatasetExtractionManifest, DatasetIdentity, ExtractionPreclassificationFailure, SourceArtifactReference, ) from gcmd_classifier.datasets.readme_discovery import ( discover_readme_candidates, select_readme_candidate, ) from gcmd_classifier.datasets.schemas import schema_for from tests.datasets.fixture_manifest import load_fixture_manifest NOW = datetime(2026, 8, 13, 16, 0, tzinfo=UTC) PUBLIC = "93.184.216.34" URL = "https://docs.example.test/readme.pdf" def _pdf_bytes( pages: list[list[tuple[float, float, str]]], *, encrypted: bool = False, image_only: bool = False, ) -> bytes: output = io.BytesIO() encryption = StandardEncryption("secret") if encrypted else None writer = canvas.Canvas(output, encrypt=encryption) for entries in pages: if image_only: image = Image.new("RGB", (20, 20), "black") writer.drawImage(ImageReader(image), 72, 700, 20, 20) for x, y, text in entries: writer.drawString(x, y, text) writer.showPage() writer.save() return output.getvalue() def _validated(pdf_bytes: bytes, tmp_path: Path, *, url: str = URL) -> RetrievedREADME: identity = DatasetIdentity( concept_id="C1-P", native_id="TARGET_001", short_name="TARGET", version="001", cmr_revision_id=3, ) view = BlindCollectionView( identity=identity, derived_native_id=identity.native_id, related_urls=({"Subtype": "READ-ME", "URL": url},), blind_view_sha256="a" * 64, ) discovery = discover_readme_candidates(view, cmr_source_sha256="b" * 64) selection = select_readme_candidate(discovery, [(0, url)], utc_now=lambda: NOW) def handler(request: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={"content-type": "application/pdf"}, content=pdf_bytes, request=request ) with httpx.Client(transport=httpx.MockTransport(handler), follow_redirects=False) as client: return retrieve_selected_readme( selection, client=client, resolver=lambda host: [PUBLIC], artifact_directory=tmp_path, utc_now=lambda: NOW, ) def _success(result: object) -> DatasetExtractionManifest: assert isinstance(result, DatasetExtractionManifest), result return result def _failure(result: object, code: str) -> ExtractionPreclassificationFailure: assert isinstance(result, ExtractionPreclassificationFailure), result assert result.error.code == code assert result.classification_attempted is False assert result.model_calls == 0 assert result.usable_for_classification is False return result def test_one_page_exact_provenance_stability_and_reconstruction(tmp_path: Path) -> None: source = _validated( _pdf_bytes([[(72, 740, "SCIENCE OVERVIEW"), (72, 720, "Exact source text.")]]), tmp_path ) first = _success(extract_pdf_manifest(source)) second = _success(extract_pdf_manifest(source)) assert first.model_dump(mode="json") == second.model_dump(mode="json") assert first.document_id == f"sha256:{source.record.artifact.sha256}" assert first.pages[0].page_number == 1 and first.pages[0].source_order == 0 assert first.pages[0].text_sha256 == hashlib.sha256(first.pages[0].text.encode()).hexdigest() assert "".join(block.text for block in first.blocks) == first.pages[0].text assert first.sections[0].label_source == "literal" heading = first.sections[0] assert ( first.pages[0].text[heading.heading_character_start : heading.heading_character_end] == heading.heading ) assert first.excluded_separators == () assert "ScienceKeywords" not in first.model_dump_json() def test_multiple_pages_repeated_headings_boundary_and_cross_page_section(tmp_path: Path) -> None: data = _pdf_bytes( [ [(72, 740, "1. OVERVIEW"), (72, 720, "Page one body")], [(72, 740, "Continuation without heading")], [(72, 740, "1. OVERVIEW"), (72, 720, "Repeated heading body")], ] ) manifest = _success(extract_pdf_manifest(_validated(data, tmp_path))) assert [page.page_number for page in manifest.pages] == [1, 2, 3] literal = [section for section in manifest.sections if section.label_source == "literal"] assert len(literal) == 2 and literal[0].section_id != literal[1].section_id assert literal[0].page_start == 1 and literal[0].page_end == 2 assert [block.source_order for block in manifest.blocks] == list(range(len(manifest.blocks))) for page in manifest.pages: assert ( "".join(block.text for block in manifest.blocks if block.page_start == page.page_number) == page.text ) def test_no_heading_uses_derived_label_and_long_text_splits_deterministically( tmp_path: Path, ) -> None: text = "word " * 200 source = _validated(_pdf_bytes([[(72, 740, text)]]), tmp_path) settings = DatasetExtractionSettings(max_block_characters=80) manifest = _success(extract_pdf_manifest(source, settings=settings)) assert manifest.sections[0].label_source == "derived" assert manifest.sections[0].heading_character_start is None assert len(manifest.blocks) > 1 assert all(len(block.text) <= 80 for block in manifest.blocks) assert "".join(block.text for block in manifest.blocks) == manifest.pages[0].text def test_blank_page_warning_and_fully_blank_failure(tmp_path: Path) -> None: mixed = _pdf_bytes([[], [(72, 740, "Text page")]]) manifest = _success(extract_pdf_manifest(_validated(mixed, tmp_path / "mixed"))) assert manifest.pages[0].extraction_status == "blank" assert any(warning.code == "BLANK_PAGE" for warning in manifest.pages[0].warnings) blank = _pdf_bytes([[]]) _failure(extract_pdf_manifest(_validated(blank, tmp_path / "blank")), "PDF_NO_EXTRACTABLE_TEXT") def test_image_only_page_and_document(tmp_path: Path) -> None: image_pdf = _pdf_bytes([[]], image_only=True) _failure(extract_pdf_manifest(_validated(image_pdf, tmp_path)), "PDF_IMAGE_ONLY") def test_encrypted_and_malformed_pdf_fail_before_inference(tmp_path: Path) -> None: encrypted = _pdf_bytes([[(72, 740, "Secret")]], encrypted=True) _failure(extract_pdf_manifest(_validated(encrypted, tmp_path / "encrypted")), "PDF_ENCRYPTED") valid = _validated(_pdf_bytes([[(72, 740, "Valid")]]), tmp_path / "malformed") malformed = b"%PDF-1.7\nnot readable\n1 0 obj\nstartxref\n0\n%%EOF" digest = hashlib.sha256(malformed).hexdigest() artifact = SourceArtifactReference( artifact_id=f"readme-{digest}", sha256=digest, byte_size=len(malformed), media_type="application/pdf", storage_reference="test/document.pdf", ) source = RetrievedREADME( record=valid.record.model_copy(update={"artifact": artifact}), content=malformed ) _failure(extract_pdf_manifest(source), "PDF_MALFORMED") def test_hash_page_character_block_and_timeout_limits(tmp_path: Path) -> None: source = _validated(_pdf_bytes([[(72, 740, "Some source text")]]), tmp_path) changed = RetrievedREADME(record=source.record, content=source.content + b"x") _failure(extract_pdf_manifest(changed), "PDF_HASH_MISMATCH") _failure( extract_pdf_manifest( source, settings=DatasetExtractionSettings(max_pages=0 + 1), pdf_opener=_two_page_fake ), "PDF_PAGE_LIMIT_EXCEEDED", ) _failure( extract_pdf_manifest(source, settings=DatasetExtractionSettings(max_characters=2)), "PDF_CHARACTER_LIMIT_EXCEEDED", ) _failure( extract_pdf_manifest( source, settings=DatasetExtractionSettings(max_blocks=1, max_block_characters=2) ), "PDF_BLOCK_LIMIT_EXCEEDED", ) clock = iter((0.0, 61.0)) _failure(extract_pdf_manifest(source, monotonic=lambda: next(clock)), "PDF_EXTRACTION_TIMEOUT") class _FakePage: width = 600.0 height = 800.0 images = [] lines = [] rects = [] def __init__(self, text: object = "text", *, words_error: bool = False) -> None: self.text = text self.words_error = words_error def extract_text(self, **kwargs): if isinstance(self.text, Exception): raise self.text return self.text def extract_words(self, **kwargs): if self.words_error: raise RuntimeError("coordinates unavailable") return [] @contextmanager def _fake_document(pages): yield type("FakePDF", (), {"pages": pages})() def _two_page_fake(stream): return _fake_document([_FakePage("one"), _FakePage("two")]) def test_page_failure_partial_failure_and_unavailable_coordinates(tmp_path: Path) -> None: source = _validated(_pdf_bytes([[(72, 740, "source")]]), tmp_path) _failure( extract_pdf_manifest( source, pdf_opener=lambda stream: _fake_document([_FakePage(RuntimeError("page"))]) ), "PDF_PAGE_EXTRACTION_FAILED", ) _failure( extract_pdf_manifest(source, pdf_opener=lambda stream: _fake_document([_FakePage(123)])), "PDF_PARTIAL_PAGE_FAILURE", ) manifest = _success( extract_pdf_manifest( source, pdf_opener=lambda stream: _fake_document([_FakePage("text", words_error=True)]) ) ) assert any(warning.code == "LAYOUT_COORDINATES_UNAVAILABLE" for warning in manifest.warnings) def test_table_multicolumn_and_low_coverage_warnings(tmp_path: Path) -> None: table = _pdf_bytes([[(72, 740, "tiny")]]) source = _validated(table, tmp_path) page = _FakePage("tiny") page.lines = [{}] * 8 def words(**kwargs): return [{"x0": x, "x1": x + 5} for x in [10, 20, 30, 40, 50, 450, 460, 470, 480, 490]] page.extract_words = words manifest = _success( extract_pdf_manifest(source, pdf_opener=lambda stream: _fake_document([page])) ) codes = {warning.code for warning in manifest.warnings} assert { "LOW_TEXT_COVERAGE", "TABLE_LIKE_LAYOUT", "LIKELY_MULTI_COLUMN", "SUSPICIOUS_READING_ORDER", } <= codes def test_policy_version_change_is_explicit_but_content_ids_remain_source_based( tmp_path: Path, ) -> None: source = _validated(_pdf_bytes([[(72, 740, "Body text")]]), tmp_path) first = _success(extract_pdf_manifest(source)) second = _success( extract_pdf_manifest( source, settings=DatasetExtractionSettings(heading_policy_version="deterministic-heading-v2"), ) ) assert first.document_id == second.document_id assert [page.page_id for page in first.pages] == [page.page_id for page in second.pages] assert first.heading_policy_version != second.heading_policy_version def test_all_frozen_pdfs_extract_deterministically_and_validate_schema(tmp_path: Path) -> None: _, cases = load_fixture_manifest(Path("tests/fixtures/datasets/cases.json")) before = { path: path.read_bytes() for case in cases for path in (case.cmr_path, case.readme_path) } outcomes = {} for case in cases: meta = case.metadata cmr_bytes = case.cmr_path.read_bytes() pdf_bytes = case.readme_path.read_bytes() validated = validate_cmr_url(meta.cmr_url) cmr_source = CMRRetrievedSource( submitted_url=meta.cmr_url, final_url=meta.cmr_url, retrieved_at="2026-08-13T00:00:00Z", status_code=200, response_bytes=cmr_bytes, source_text=cmr_bytes.decode(), sha256=meta.cmr_sha256, parsed_json=json.loads(cmr_bytes), ) resolved = resolve_cmr_collection(validated, cmr_source) discovery = discover_readme_candidates( resolved.blind_view, cmr_source_sha256=meta.cmr_sha256 ) candidate = next( item for item in discovery.candidates if item.url == meta.selected_readme_url ) selection = select_readme_candidate( discovery, [(candidate.source_index, candidate.url)], utc_now=lambda: NOW ) def handler(request, expected=meta.selected_readme_url, content=pdf_bytes): assert str(request.url) == expected return httpx.Response( 200, headers={"content-type": "application/pdf"}, content=content, request=request ) with httpx.Client(transport=httpx.MockTransport(handler), follow_redirects=False) as client: retrieved = retrieve_selected_readme( selection, client=client, resolver=lambda host: [PUBLIC], artifact_directory=tmp_path / meta.case_id, utc_now=lambda: NOW, ) first = extract_pdf_manifest(retrieved) second = extract_pdf_manifest(retrieved) manifest = _success(first) assert manifest.model_dump(mode="json") == _success(second).model_dump(mode="json") assert manifest.document.sha256 == meta.readme_sha256 assert manifest.lineage.cmr_source_sha256 == meta.cmr_sha256 assert manifest.lineage.identity.concept_id == meta.expected_concept_id assert "ScienceKeywords" not in manifest.model_dump_json() DatasetExtractionManifest.model_validate(manifest.model_dump(mode="json")) assert schema_for(DatasetExtractionManifest)["title"] == "DatasetExtractionManifest" outcomes[meta.case_id] = (manifest.page_count, manifest.block_count, len(manifest.warnings)) assert {path: path.read_bytes() for path in before} == before assert len(outcomes) == 4 def test_multiple_headings_on_page_and_heading_at_page_end(tmp_path: Path) -> None: source = _validated( _pdf_bytes( [ [ (72, 740, "1. FIRST HEADING"), (72, 720, "First body"), (72, 700, "2. END HEADING"), ], [(72, 740, "Body continuing from end heading")], ] ), tmp_path, ) manifest = _success(extract_pdf_manifest(source)) literal = [section for section in manifest.sections if section.label_source == "literal"] assert [section.heading for section in literal] == ["1. FIRST HEADING", "2. END HEADING"] assert literal[-1].page_start == 1 and literal[-1].page_end == 2 assert len(set(section.section_id for section in literal)) == 2