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, )