wehub-resource-sync
e4dcfc49aa
Tests / Import Check (Python 3.13) (push) Has been cancelled
Tests / Import Check (Python 3.14) (push) Has been cancelled
Tests / Python Tests (Python 3.11) (push) Has been cancelled
Tests / Python Tests (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.14) (push) Has been cancelled
Tests / Test Summary (push) Has been cancelled
Tests / Lint and Format (push) Has been cancelled
Tests / Web Node Tests (push) Has been cancelled
Tests / Import Check (Python 3.11) (push) Has been cancelled
Tests / Import Check (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.13) (push) Has been cancelled
377 lines
13 KiB
Python
377 lines
13 KiB
Python
"""Tests for deeptutor.utils.document_extractor."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
|
|
from docx import Document as DocxDocument
|
|
from openpyxl import Workbook
|
|
from pptx import Presentation
|
|
from pptx.util import Inches
|
|
import pytest
|
|
|
|
from deeptutor.utils import document_extractor as document_extractor_module
|
|
from deeptutor.utils.document_extractor import (
|
|
MAX_DOC_BYTES,
|
|
MAX_EXTRACTED_CHARS_PER_DOC,
|
|
CorruptDocumentError,
|
|
DocumentTooLargeError,
|
|
EmptyDocumentError,
|
|
UnsupportedDocumentError,
|
|
extract_documents_from_records,
|
|
extract_text_from_bytes,
|
|
extract_text_from_path,
|
|
is_document_extension,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures — generate office docs on the fly
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_docx(paragraphs: list[str]) -> bytes:
|
|
doc = DocxDocument()
|
|
for p in paragraphs:
|
|
doc.add_paragraph(p)
|
|
buf = io.BytesIO()
|
|
doc.save(buf)
|
|
return buf.getvalue()
|
|
|
|
|
|
def _make_xlsx(sheets: dict[str, list[list[object]]]) -> bytes:
|
|
wb = Workbook()
|
|
default = wb.active
|
|
first = True
|
|
for name, rows in sheets.items():
|
|
ws = default if first else wb.create_sheet()
|
|
ws.title = name
|
|
for row in rows:
|
|
ws.append(row)
|
|
first = False
|
|
buf = io.BytesIO()
|
|
wb.save(buf)
|
|
return buf.getvalue()
|
|
|
|
|
|
def _make_pptx(slides_text: list[list[str]]) -> bytes:
|
|
prs = Presentation()
|
|
for slide_texts in slides_text:
|
|
slide = prs.slides.add_slide(prs.slide_layouts[5]) # blank-ish layout
|
|
for i, text in enumerate(slide_texts):
|
|
tb = slide.shapes.add_textbox(Inches(1), Inches(1 + i * 0.5), Inches(6), Inches(0.5))
|
|
tb.text_frame.text = text
|
|
buf = io.BytesIO()
|
|
prs.save(buf)
|
|
return buf.getvalue()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# is_document_extension
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestIsDocumentExtension:
|
|
def test_office(self) -> None:
|
|
assert is_document_extension("foo.pdf")
|
|
assert is_document_extension("foo.DOCX")
|
|
assert is_document_extension("report.xlsx")
|
|
assert is_document_extension("deck.pptx")
|
|
|
|
def test_text_and_code(self) -> None:
|
|
# Any extension in FileTypeRouter.TEXT_EXTENSIONS should be supported.
|
|
assert is_document_extension("notes.txt")
|
|
assert is_document_extension("readme.md")
|
|
assert is_document_extension("module.py")
|
|
assert is_document_extension("config.yaml")
|
|
assert is_document_extension("data.json")
|
|
assert is_document_extension("index.html")
|
|
assert is_document_extension("table.csv")
|
|
|
|
def test_unsupported(self) -> None:
|
|
assert not is_document_extension("foo.png")
|
|
assert not is_document_extension("foo.zip")
|
|
assert not is_document_extension("foo.exe")
|
|
assert not is_document_extension("foo")
|
|
assert not is_document_extension("")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# extract_text_from_bytes — happy paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractDocx:
|
|
def test_basic_paragraphs(self) -> None:
|
|
data = _make_docx(["Hello world", "Second paragraph", ""])
|
|
text = extract_text_from_bytes("doc.docx", data)
|
|
assert "Hello world" in text
|
|
assert "Second paragraph" in text
|
|
|
|
def test_path_helper_can_disable_chat_truncation(self, tmp_path) -> None:
|
|
data = _make_docx(["a" * (MAX_EXTRACTED_CHARS_PER_DOC + 10)])
|
|
path = tmp_path / "long.docx"
|
|
path.write_bytes(data)
|
|
|
|
text = extract_text_from_path(path, max_chars=None)
|
|
|
|
assert len(text) > MAX_EXTRACTED_CHARS_PER_DOC
|
|
assert "truncated" not in text
|
|
|
|
def test_ooxml_fallback_without_python_docx(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(document_extractor_module, "DocxDocument", None)
|
|
data = _make_docx(["Fallback paragraph", "第二段"])
|
|
|
|
text = extract_text_from_bytes("doc.docx", data)
|
|
|
|
assert "Fallback paragraph" in text
|
|
assert "第二段" in text
|
|
|
|
|
|
class TestExtractXlsx:
|
|
def test_multiple_sheets(self) -> None:
|
|
data = _make_xlsx(
|
|
{
|
|
"Alpha": [["a1", "b1"], ["a2", 42]],
|
|
"Beta": [["x", "y"]],
|
|
}
|
|
)
|
|
text = extract_text_from_bytes("book.xlsx", data)
|
|
assert "--- Sheet: Alpha ---" in text
|
|
assert "--- Sheet: Beta ---" in text
|
|
assert "a1" in text and "42" in text
|
|
assert "x" in text
|
|
|
|
def test_ooxml_fallback_without_openpyxl(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(document_extractor_module, "load_workbook", None)
|
|
data = _make_xlsx({"Alpha": [["name", "score"], ["alice", 98]]})
|
|
|
|
text = extract_text_from_bytes("book.xlsx", data)
|
|
|
|
assert "--- Sheet: Alpha ---" in text
|
|
assert "alice" in text
|
|
assert "98" in text
|
|
|
|
|
|
class TestExtractPptx:
|
|
def test_basic_slides(self) -> None:
|
|
data = _make_pptx([["Slide 1 title", "Slide 1 body"], ["Slide 2 only text"]])
|
|
text = extract_text_from_bytes("deck.pptx", data)
|
|
assert "--- Slide 1 ---" in text
|
|
assert "--- Slide 2 ---" in text
|
|
assert "Slide 1 title" in text
|
|
assert "Slide 2 only text" in text
|
|
|
|
def test_ooxml_fallback_without_python_pptx(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(document_extractor_module, "PptxPresentation", None)
|
|
data = _make_pptx([["Fallback slide", "第二行"]])
|
|
|
|
text = extract_text_from_bytes("deck.pptx", data)
|
|
|
|
assert "--- Slide 1 ---" in text
|
|
assert "Fallback slide" in text
|
|
assert "第二行" in text
|
|
|
|
|
|
class TestExtractTextLike:
|
|
def test_plain_txt(self) -> None:
|
|
text = extract_text_from_bytes("note.txt", "hello world\nline two".encode("utf-8"))
|
|
assert "hello world" in text
|
|
assert "line two" in text
|
|
|
|
def test_python_source(self) -> None:
|
|
src = b"def greet(name: str) -> str:\n return f'hi {name}'\n"
|
|
text = extract_text_from_bytes("greet.py", src)
|
|
assert "def greet" in text
|
|
|
|
def test_json(self) -> None:
|
|
text = extract_text_from_bytes("data.json", b'{"x": 42, "y": "ok"}')
|
|
assert '"x": 42' in text
|
|
|
|
def test_csv(self) -> None:
|
|
text = extract_text_from_bytes("table.csv", b"a,b,c\n1,2,3\n")
|
|
assert "a,b,c" in text
|
|
assert "1,2,3" in text
|
|
|
|
def test_markdown(self) -> None:
|
|
text = extract_text_from_bytes("doc.md", b"# Heading\n\nBody.\n")
|
|
assert "# Heading" in text
|
|
|
|
def test_utf8_with_bom(self) -> None:
|
|
# The candidate chain has utf-8 before utf-8-sig (same order as KB
|
|
# pipeline), so BOM-prefixed bytes decode as utf-8 and the BOM is
|
|
# retained. Mirror that behavior here.
|
|
text = extract_text_from_bytes("note.txt", b"\xef\xbb\xbfhello")
|
|
assert "hello" in text
|
|
|
|
def test_gbk_fallback(self) -> None:
|
|
# "你好" in GBK
|
|
text = extract_text_from_bytes("note.txt", "你好".encode("gbk"))
|
|
assert text == "你好"
|
|
|
|
def test_svg(self) -> None:
|
|
svg = (
|
|
b'<?xml version="1.0"?>'
|
|
b'<svg xmlns="http://www.w3.org/2000/svg" width="100" height="100">'
|
|
b'<circle cx="50" cy="50" r="40" fill="red"/>'
|
|
b'<text x="50" y="55">Hello</text>'
|
|
b"</svg>"
|
|
)
|
|
text = extract_text_from_bytes("logo.svg", svg)
|
|
assert "<svg" in text
|
|
assert "<circle" in text
|
|
assert "Hello" in text
|
|
|
|
|
|
class TestExtractPdf:
|
|
def test_minimal_pdf(self) -> None:
|
|
# Build a minimal valid PDF via pymupdf (dependency already in project).
|
|
pytest.importorskip("fitz")
|
|
import fitz # noqa: WPS433
|
|
|
|
doc = fitz.open()
|
|
page = doc.new_page()
|
|
page.insert_text((72, 72), "Hello PDF world")
|
|
buf = io.BytesIO()
|
|
doc.save(buf)
|
|
doc.close()
|
|
data = buf.getvalue()
|
|
|
|
text = extract_text_from_bytes("sample.pdf", data)
|
|
assert "Hello PDF world" in text
|
|
assert "--- Page 1 ---" in text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Failure modes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFailureModes:
|
|
def test_unsupported_extension(self) -> None:
|
|
with pytest.raises(UnsupportedDocumentError):
|
|
extract_text_from_bytes("foo.zip", b"\x00\x00")
|
|
|
|
def test_empty_bytes(self) -> None:
|
|
with pytest.raises(EmptyDocumentError):
|
|
extract_text_from_bytes("foo.docx", b"")
|
|
|
|
def test_too_large(self) -> None:
|
|
fake = b"PK\x03\x04" + b"\x00" * (MAX_DOC_BYTES + 1)
|
|
with pytest.raises(DocumentTooLargeError):
|
|
extract_text_from_bytes("foo.docx", fake)
|
|
|
|
def test_pdf_magic_mismatch(self) -> None:
|
|
with pytest.raises(CorruptDocumentError):
|
|
extract_text_from_bytes("foo.pdf", b"this is not a pdf")
|
|
|
|
def test_ooxml_magic_mismatch(self) -> None:
|
|
with pytest.raises(CorruptDocumentError):
|
|
extract_text_from_bytes("foo.docx", b"not an office file")
|
|
|
|
def test_corrupt_docx(self) -> None:
|
|
# OOXML header but garbage body
|
|
with pytest.raises(CorruptDocumentError):
|
|
extract_text_from_bytes("foo.docx", b"PK\x03\x04" + b"\x00" * 512)
|
|
|
|
def test_empty_docx_no_text(self) -> None:
|
|
data = _make_docx([]) # no paragraphs
|
|
with pytest.raises(EmptyDocumentError):
|
|
extract_text_from_bytes("foo.docx", data)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Truncation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTruncation:
|
|
def test_long_docx_is_truncated(self) -> None:
|
|
# single paragraph of 250k chars → well over the 200k per-doc cap
|
|
long_text = "a" * 250_000
|
|
data = _make_docx([long_text])
|
|
text = extract_text_from_bytes("big.docx", data)
|
|
assert len(text) <= MAX_EXTRACTED_CHARS_PER_DOC + 200 # allow notice suffix
|
|
assert "truncated" in text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# extract_documents_from_records
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractDocumentsFromRecords:
|
|
def test_mixed_image_and_doc(self) -> None:
|
|
docx_bytes = _make_docx(["hello there"])
|
|
docx_b64 = base64.b64encode(docx_bytes).decode()
|
|
image_b64 = base64.b64encode(b"\x89PNG\r\n\x1a\n").decode()
|
|
|
|
records = [
|
|
{
|
|
"type": "image",
|
|
"filename": "pic.png",
|
|
"base64": image_b64,
|
|
"mime_type": "image/png",
|
|
"url": "",
|
|
},
|
|
{
|
|
"type": "file",
|
|
"filename": "note.docx",
|
|
"base64": docx_b64,
|
|
"mime_type": "",
|
|
"url": "",
|
|
},
|
|
]
|
|
|
|
doc_texts, updated = extract_documents_from_records(records)
|
|
|
|
assert len(doc_texts) == 1
|
|
assert "[File: note.docx]" in doc_texts[0]
|
|
assert "hello there" in doc_texts[0]
|
|
|
|
# image record untouched
|
|
assert updated[0]["base64"] == image_b64
|
|
# doc record base64 cleared, extracted_chars set
|
|
assert updated[1]["base64"] == ""
|
|
assert updated[1]["extracted_chars"] > 0
|
|
|
|
def test_unsupported_record_is_passthrough(self) -> None:
|
|
records = [
|
|
{"type": "file", "filename": "foo.zip", "base64": "AAAA", "mime_type": "", "url": ""}
|
|
]
|
|
doc_texts, updated = extract_documents_from_records(records)
|
|
assert doc_texts == []
|
|
assert updated[0]["base64"] == "AAAA" # untouched — not a doc extension
|
|
|
|
def test_failed_extraction_emits_error_marker(self) -> None:
|
|
records = [
|
|
{
|
|
"type": "file",
|
|
"filename": "bad.pdf",
|
|
"base64": base64.b64encode(b"not a pdf").decode(),
|
|
"mime_type": "",
|
|
"url": "",
|
|
}
|
|
]
|
|
doc_texts, updated = extract_documents_from_records(records)
|
|
assert len(doc_texts) == 1
|
|
assert "bad.pdf" in doc_texts[0]
|
|
assert "could not be read" in doc_texts[0]
|
|
assert updated[0]["base64"] == "" # stripped even on failure
|
|
|
|
def test_invalid_base64_emits_error_marker(self) -> None:
|
|
records = [
|
|
{
|
|
"type": "file",
|
|
"filename": "bad.docx",
|
|
"base64": "!!!not base64!!!",
|
|
"mime_type": "",
|
|
"url": "",
|
|
}
|
|
]
|
|
doc_texts, updated = extract_documents_from_records(records)
|
|
# invalid base64 with validate=False may silently decode or emit error — both
|
|
# paths end up as an error marker since resulting bytes won't pass magic check
|
|
assert len(doc_texts) == 1
|
|
assert "bad.docx" in doc_texts[0]
|