DataOps-platform/tests/data_research/test_document_extractors.py

from __future__ import annotations

import io

from docx import Document

from app.core.data_research.extractors.base import ExtractionContext


def docx_bytes(*, empty=False):
    document = Document()
    if not empty:
        document.add_paragraph("客户主数据定义")
        table = document.add_table(rows=3, cols=2)
        table.rows[0].cells[0].text = "字段"
        table.rows[0].cells[1].text = "定义"
        table.rows[1].cells[0].text = "customer_id"
        table.rows[1].cells[1].text = "客户唯一编号"
        table.rows[2].cells[0].text = "字段"
        table.rows[2].cells[1].text = "定义"
    stream = io.BytesIO()
    document.save(stream)
    return stream.getvalue()


def test_docx_extracts_paragraphs_and_table_cells_with_exact_locations():
    from app.core.data_research.extractors.docx import DocxExtractor

    batch = DocxExtractor().extract(
        docx_bytes(),
        ExtractionContext(
            filename="governance.docx",
            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ),
    )

    assert batch.parser_version == "docx-v1"
    assert batch.items[0].data["text"] == "客户主数据定义"
    assert batch.items[0].evidence[0].locator == {"kind": "docx.paragraph", "paragraph": 1}
    table_cells = [item for item in batch.items if item.data["kind"] == "table_cell"]
    assert table_cells[-1].data["text"] == "客户唯一编号"
    assert table_cells[-1].evidence[0].locator == {
        "kind": "docx.table",
        "table": 1,
        "row": 2,
        "cell": 2,
    }
    assert len(table_cells) == 4  # repeated header row is removed


def test_empty_docx_returns_a_deterministic_warning():
    from app.core.data_research.extractors.docx import DocxExtractor

    batch = DocxExtractor().extract(
        docx_bytes(empty=True),
        ExtractionContext(filename="empty.docx", media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document"),
    )
    assert batch.items == ()
    assert batch.warnings == ("document contains no usable text",)


class FakePage:
    def __init__(self, text, tables=(), bbox=(0, 0, 600, 800)):
        self._text = text
        self._tables = tables
        self.bbox = bbox

    def extract_text(self):
        return self._text

    def extract_tables(self):
        return list(self._tables)


class FakePdf:
    def __init__(self, pages):
        self.pages = pages

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return False


def test_pdf_preserves_page_order_tables_and_marks_scanned_pages(monkeypatch):
    import pdfplumber

    from app.core.data_research.extractors.pdf import PdfExtractor

    pages = [
        FakePage("第一页定义", tables=[[['字段', '定义'], ['id', '唯一编号']]]),
        FakePage(None),
        FakePage("第三页定义"),
    ]
    monkeypatch.setattr(pdfplumber, "open", lambda _stream: FakePdf(pages))

    batch = PdfExtractor().extract(
        b"%PDF-fixture",
        ExtractionContext(filename="standard.pdf", media_type="application/pdf"),
    )

    assert [item.evidence[0].locator["page"] for item in batch.items] == [1, 1, 1, 1, 1, 3]
    assert batch.items[0].evidence[0].locator == {
        "kind": "pdf.page",
        "page": 1,
        "bbox": [0.0, 0.0, 1.0, 1.0],
    }
    assert batch.metadata["ocr_required_pages"] == [2]
    assert batch.warnings == ("OCR required for pages: 2",)


def test_registry_resolves_docx_and_pdf():
    from app.core.data_research.extractors.docx import DocxExtractor
    from app.core.data_research.extractors.pdf import PdfExtractor
    from app.core.data_research.extractors.registry import default_registry

    registry = default_registry()
    assert isinstance(registry.resolve("application/pdf", "a.pdf"), PdfExtractor)
    assert isinstance(
        registry.resolve(
            "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            "a.docx",
        ),
        DocxExtractor,
    )