chore: import upstream snapshot with attribution

tests/utils/test_document_extractor.py

@@ -0,0 +1,376 @@
+"""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]