"""Unit tests for the extract stage. Run: python3 -m unittest discover . Every fixture is generated here with the stdlib (zip + XML for office formats, hand-built PDFs), so the suite needs no binary blobs and runs on any OS. External tools are mocked; the real-tool path is exercised only where the machine happens to have them. """ import io import json import os import shutil import tempfile import unittest import zipfile import zlib from unittest import mock from . import extract as ex from .manifest import validate_document from .walk import walk def manifest(root, **scope): return validate_document({ "freeagent": 1, "id": "t", "owner": {"host": "h"}, "scope": dict({"roots": [root]}, **scope), "mask": ["read"], "stages": ["walk", "extract"], "sweep": "manual", "sink": {"catalog": "local", "vectors": "local", "plaintext": "never"}, }) def put(root, rel, data): path = os.path.join(root, *rel.split("/")) os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "wb") as fh: fh.write(data if isinstance(data, bytes) else data.encode("utf-8")) def zipped(members): buf = io.BytesIO() with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: for name, data in members.items(): zf.writestr(name, data) return buf.getvalue() DOCX_XML = ('' 'HelloWorld' 'Second paragraph') XLSX_SST = ('' 'NameAda') XLSX_SHEET = ('' '042' '1inline' '') PPTX_SLIDE = ('' 'Slide ' 'title') ODT_CONTENT = ('' 'HeadingBody span' 'c1' 'c2' '') def pdf_bytes(content, compress=False): stream = zlib.compress(content) if compress else content filt = b"/Filter /FlateDecode " if compress else b"" return (b"%PDF-1.4\n1 0 obj << /Type /Catalog >> endobj\n" b"2 0 obj << " + filt + b"/Length " + str(len(stream)).encode() + b" >>\nstream\n" + stream + b"\nendstream\nendobj\n%%EOF\n") class ExtractorUnitTests(unittest.TestCase): def test_decode_bytes(self): self.assertEqual(ex.decode_bytes("héllo".encode("utf-8")), ("héllo", "utf-8")) self.assertEqual(ex.decode_bytes(b"\xef\xbb\xbfhi"), ("hi", "utf-8-sig")) self.assertEqual(ex.decode_bytes("hi".encode("utf-16"))[0], "hi") self.assertEqual(ex.decode_bytes("hi there".encode("utf-16-le")), ("hi there", "utf-16-le")) self.assertEqual(ex.decode_bytes(b"caf\xe9"), ("café", "cp1252")) def test_html_to_text(self): out = ex.html_to_text("T" "

One & two

Three
") self.assertIn("One & two", out) self.assertIn("Three", out) self.assertNotIn("bad()", out) self.assertNotIn("x{}", out) def test_docx(self): text = ex.extract_docx(zipped({"word/document.xml": DOCX_XML})) self.assertEqual(text.strip().splitlines(), ["Hello\tWorld", "Second paragraph"]) def test_xlsx(self): text = ex.extract_xlsx(zipped({"xl/sharedStrings.xml": XLSX_SST, "xl/worksheets/sheet1.xml": XLSX_SHEET})) self.assertIn("Name\t42", text) self.assertIn("Ada\tinline", text) def test_pptx(self): text = ex.extract_pptx(zipped({"ppt/slides/slide2.xml": PPTX_SLIDE, "ppt/slides/slide1.xml": PPTX_SLIDE.replace("title", "first")})) lines = [l for l in text.splitlines() if l] self.assertEqual(lines, ["# slide1", "Slide first", "# slide2", "Slide title"]) def test_odf(self): text = ex.extract_odf(zipped({"content.xml": ODT_CONTENT})) self.assertEqual(text.splitlines(), ["Heading", "Body span", "c1\tc2"]) def test_zip_member_cap(self): big = zipped({"content.xml": b"x" * (ex.ZIP_MEMBER_CAP + 1)}) with self.assertRaises(ValueError): ex.extract_odf(big) def test_office_refuses_a_dtd(self): # Office XML never carries a DTD; one is how entity expansion gets in. bomb = (b'' b']>' b'' b'&b;') with self.assertRaises(ValueError): ex.extract_docx(zipped({"word/document.xml": bomb})) def test_office_text_cap_stops_with_the_text_so_far(self): rows = "".join('row %d' % i for i in range(200)) sheet = ('' "%s" % rows).encode() with mock.patch.object(ex, "OFFICE_TEXT_CAP", 300): text = ex.extract_xlsx(zipped({"xl/worksheets/sheet1.xml": sheet})) self.assertIn("row 0", text) self.assertNotIn("row 199", text) self.assertTrue(text.splitlines()[-1].startswith("[stopped reading here:")) def test_office_decompressed_cap_counts_every_member(self): para = b"x" * 2000 doc = (b'' b"" + para + b"") files = {"word/document.xml": doc} files.update({"word/header%d.xml" % i: doc for i in range(1, 6)}) with mock.patch.object(ex, "OFFICE_DOC_CAP", len(doc) * 3): text = ex.extract_docx(zipped(files)) self.assertIn("[stopped reading here: the document expands past", text) def test_office_parse_is_streamed(self): # Rows are dropped as they are read: the parsed tree never holds them all. import tracemalloc rows = "".join('%s' % ("v" * 50) for _ in range(60000)) sheet = ('' "%s" % rows).encode() data = zipped({"xl/worksheets/sheet1.xml": sheet}) tracemalloc.start() text = ex.extract_xlsx(data) peak = tracemalloc.get_traced_memory()[1] tracemalloc.stop() self.assertEqual(text.count("v" * 50), 60000) # the XML is ~4.6 MB; a whole-tree parse peaks at many times that self.assertLess(peak, 3 * len(sheet), peak) def test_office_drops_unwanted_siblings_linear(self): body = b"hi" * 100000 doc = (b'' b"" + body + b"") import time as _t t0 = _t.monotonic() text = ex.extract_docx(zipped({"word/document.xml": doc})) self.assertLess(_t.monotonic() - t0, 10) self.assertEqual(text.count("hi"), 100000) self.assertNotIn("[stopped", text) def test_office_one_paragraph_with_endless_parts_is_cut_and_reading_goes_on(self): runs = b"a" * 10 doc = (b'' b"" + runs + b"after" b"") with mock.patch.object(ex, "OFFICE_PARTS_PER_ITEM", 15): text = ex.extract_docx(zipped({"word/document.xml": doc})) lines = text.splitlines() self.assertTrue(lines[0].startswith("a") and len(lines[0]) < 10, lines[0]) self.assertIn("after", lines) self.assertEqual(lines[-1], "[1 paragraph(s) or row(s) held more than 15 parts; each was cut there]") def test_office_dropping_past_the_cap_is_linear(self): # A closing element is not its parent's last child when the pull # parser has read ahead; per-element remove() was quadratic (400k # runs: 60 s and the rest of the document lost). import time as _t runs = b"a" * 60000 doc = (b'' b"" + runs + b"TAIL" b"") with mock.patch.object(ex, "OFFICE_PARTS_PER_ITEM", 20000): t0 = _t.monotonic() text = ex.extract_docx(zipped({"word/document.xml": doc})) self.assertLess(_t.monotonic() - t0, 10) self.assertIn("TAIL", text.splitlines()) def test_office_cut_row_has_no_phantom_cells(self): cells = "".join('v' for _ in range(40)) sheet = ('' "%snext" "" % cells).encode() with mock.patch.object(ex, "OFFICE_PARTS_PER_ITEM", 30): text = ex.extract_xlsx(zipped({"xl/worksheets/sheet1.xml": sheet})) first = text.splitlines()[1] self.assertTrue(all(c == "v" for c in first.split("\t")), first) self.assertIn("next", text.splitlines()) def test_office_utf16_is_read_and_its_dtd_still_refused(self): ns = 'xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"' good = ('' "wide caf\u00e9" % ns) for enc in ("utf-16", "utf-16-be"): data = good.encode(enc) if enc == "utf-16-be": data = b"\xfe\xff" + data self.assertEqual(ex.extract_docx(zipped({"word/document.xml": data})).strip(), "wide caf\u00e9") bomb = (']>' "&a;" % ns) with self.assertRaises(ValueError): ex.extract_docx(zipped({"word/document.xml": bomb.encode("utf-16")})) with self.assertRaises(ValueError): # UTF-32: not an Office encoding ex.extract_docx(zipped({"word/document.xml": good.replace("UTF-16", "UTF-32").encode("utf-32-le")})) def test_docx_text_box_follows_the_paragraph_holding_it(self): ns = b'xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"' doc = (b"" b"before" b"holder" b"boxed" b"after") text = ex.extract_docx(zipped({"word/document.xml": doc})) self.assertEqual([l for l in text.splitlines() if l], ["before", "holder", "boxed", "after"]) def test_eml_plain_and_html(self): plain = b"From: a@b.c\nTo: d@e.f\nSubject: Hi\n\nbody text\n" out = ex.extract_eml(plain) self.assertIn("Subject: Hi", out) self.assertIn("body text", out) htm = (b"From: a@b.c\nSubject: H\nContent-Type: text/html\n\n" b"

rich body

") out = ex.extract_eml(htm) self.assertIn("rich body", out) self.assertNotIn("", out) def test_rtf(self): rtf = (rb"{\rtf1\ansi{\fonttbl{\f0 Arial;}}{\colortbl;\red0;}" rb"\f0 Hello\par caf\'e9 \u8364? end\tab tabbed}") out = ex.extract_rtf(rtf) self.assertEqual(out, "Hello\ncafé € end\ttabbed") self.assertNotIn("Arial", out) def test_pdf_builtin_plain_and_flate(self): content = b"BT /F1 12 Tf 72 700 Td (Hello) Tj T* [(Wor) -20 (ld)] TJ ET" self.assertIn("Hello", ex.extract_pdf_builtin(pdf_bytes(content))) out = ex.extract_pdf_builtin(pdf_bytes(content, compress=True)) self.assertIn("Hello", out) self.assertIn("World", out.replace(" ", "").replace("\n", "") and out) def test_pdf_builtin_refuses_garbage_and_encrypted(self): garbage = b"BT (\x01\x02\x03\x04\x05\x06\x07\x08\x0b\x0e\x0f\x10\x11\x12) Tj ET" self.assertIsNone(ex.extract_pdf_builtin(pdf_bytes(garbage))) self.assertIsNone(ex.extract_pdf_builtin(pdf_bytes(b"BT (x) Tj ET") + b"/Encrypt 5 0 R")) self.assertIsNone(ex.extract_pdf_builtin(pdf_bytes(b"no text ops here"))) def test_pdf_escapes(self): from . import pdftext as pt self.assertEqual(pt._Lexer(rb"(a\)b\\c\101\n)").token(), ("str", b"a)b\\cA\n")) class StageCase(unittest.TestCase): def setUp(self): self.root = tempfile.mkdtemp(prefix="fa-ex-") self.out = tempfile.mkdtemp(prefix="fa-exout-") put(self.root, "a.txt", "plain text\n") put(self.root, "copy-of-a.txt", "plain text\n") put(self.root, "latin.txt", "caf\xe9".encode("cp1252")) put(self.root, "page.html", "

hello web

") put(self.root, "data.csv", "a,b\n1,2\n") put(self.root, "conf.json", '{"k": 1}') put(self.root, "mail.eml", "From: a@b.c\nSubject: S\n\nhi\n") put(self.root, "doc.docx", zipped({"word/document.xml": DOCX_XML})) put(self.root, "sheet.xlsx", zipped({"xl/sharedStrings.xml": XLSX_SST, "xl/worksheets/sheet1.xml": XLSX_SHEET})) put(self.root, "deck.pptx", zipped({"ppt/slides/slide1.xml": PPTX_SLIDE})) put(self.root, "text.odt", zipped({"content.xml": ODT_CONTENT})) put(self.root, "doc.pdf", pdf_bytes(b"BT (pdf words) Tj ET", compress=True)) put(self.root, "scan.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 32) put(self.root, "blob.bin", b"\x00\x01\x02\x03") put(self.root, "empty.txt", "") put(self.root, "noext-empty", "") put(self.root, "broken.docx", b"PK\x03\x04not really a zip") put(self.root, "big.txt", "x" * 100_000) def tearDown(self): shutil.rmtree(self.root, ignore_errors=True) shutil.rmtree(self.out, ignore_errors=True) def run_stage(self, use_external=False, **scope): m = manifest(self.root, **scope) walk(m, self.out) s = ex.extract(m, self.out, use_external=use_external) with open(os.path.join(self.out, "extract.jsonl"), encoding="utf-8") as fh: idx = {json.loads(l)["path"]: json.loads(l) for l in fh} return s, idx def text_of(self, rec): with open(ex.store_path(self.out, rec["sha256"]), encoding="utf-8") as fh: return fh.read() class StageTests(StageCase): def test_everything_builtin(self): s, idx = self.run_stage(max_file_bytes=50_000) self.assertEqual(idx["a.txt"]["status"], "ok") self.assertEqual(idx["a.txt"]["extractor"], "text:utf-8") self.assertEqual(self.text_of(idx["a.txt"]), "plain text\n") self.assertEqual(idx["latin.txt"]["extractor"], "text:cp1252") self.assertEqual(self.text_of(idx["latin.txt"]), "café\n") self.assertEqual(self.text_of(idx["page.html"]).strip(), "hello web") self.assertEqual(self.text_of(idx["data.csv"]), "a\tb\n1\t2\n") self.assertEqual(idx["conf.json"]["status"], "ok") self.assertIn("Subject: S", self.text_of(idx["mail.eml"])) self.assertIn("Hello\tWorld", self.text_of(idx["doc.docx"])) self.assertIn("Name\t42", self.text_of(idx["sheet.xlsx"])) self.assertIn("Slide title", self.text_of(idx["deck.pptx"])) self.assertIn("Heading", self.text_of(idx["text.odt"])) self.assertEqual(idx["doc.pdf"]["extractor"], "pdf-builtin") self.assertIn("pdf words", self.text_of(idx["doc.pdf"])) self.assertEqual(idx["scan.png"]["status"], "unsupported") self.assertIn("tesseract", idx["scan.png"]["error"]) self.assertEqual(idx["blob.bin"]["status"], "unsupported") self.assertEqual(idx["empty.txt"]["status"], "empty") # same (empty) content as empty.txt -> deduped onto its record self.assertEqual(idx["noext-empty"]["status"], "empty") self.assertEqual(idx["noext-empty"]["sha256"], idx["empty.txt"]["sha256"]) # the inode/x-empty branch on its own self.assertEqual(ex.Extractor(use_external=False).extract("/nonexistent", "inode/x-empty"), ("ok", "empty", "", None)) self.assertEqual(idx["broken.docx"]["status"], "error") self.assertEqual(idx["big.txt"]["status"], "skipped") self.assertEqual(idx["big.txt"]["error"], "too_large") # duplicate content: one store entry, both paths indexed self.assertEqual(idx["copy-of-a.txt"]["sha256"], idx["a.txt"]["sha256"]) c = s["counts"] self.assertEqual(c["files"], len(idx)) self.assertEqual(c["skipped"], 1) self.assertEqual(c["error"], 1) self.assertEqual(c["unsupported"], 2) self.assertEqual(c["empty"], 2) # empty.txt + its content-twin noext-empty self.assertEqual(s["tools"], {"pdftotext": None, "tesseract": None}) def test_reuse_across_sweeps(self): s1, _ = self.run_stage() s2, idx = self.run_stage() # everything that produced text (or was empty) last sweep, incl. the duplicate copy self.assertEqual(s2["counts"]["reused"], s1["counts"]["ok"] + s1["counts"]["empty"]) self.assertEqual(s2["counts"]["ok"], 0) self.assertTrue(idx["a.txt"]["reused"]) # unsupported verdicts are reused while the tool set is unchanged self.assertTrue(idx["scan.png"].get("reused")) # a changed file is re-extracted put(self.root, "a.txt", "new text\n") s3, idx3 = self.run_stage() self.assertEqual(self.text_of(idx3["a.txt"]), "new text\n") self.assertFalse(idx3["a.txt"].get("reused")) def test_cached_verdicts_invalidate_on_extractor_upgrade(self): self.run_stage() with mock.patch.object(ex, "EXTRACT_VERSION", ex.EXTRACT_VERSION + 1): s, idx = self.run_stage() self.assertEqual(s["counts"]["reused"], 0) self.assertFalse(idx["a.txt"].get("reused")) self.assertFalse(idx["scan.png"].get("reused")) def test_unsupported_retried_when_tools_change(self): self.run_stage() fake_tess = lambda path, tool: "OCR TEXT\n" with mock.patch.object(ex, "tool_path", lambda n: "/usr/bin/" + n), \ mock.patch.object(ex, "extract_image_tool", fake_tess), \ mock.patch.object(ex, "extract_pdf_tool", lambda p, t: "from poppler\n"): s, idx = self.run_stage(use_external=True) self.assertEqual(idx["scan.png"]["status"], "ok") self.assertEqual(idx["scan.png"]["extractor"], "tesseract") self.assertEqual(self.text_of(idx["scan.png"]), "OCR TEXT\n") self.assertEqual(s["tools"]["pdftotext"], "/usr/bin/pdftotext") # the pdf was already extracted by the builtin path and is reused, not re-run self.assertTrue(idx["doc.pdf"].get("reused")) def test_tool_failure_is_an_error_entry(self): def boom(path, tool): raise RuntimeError("tesseract exited 1: no language data") with mock.patch.object(ex, "tool_path", lambda n: "/usr/bin/" + n if n == "tesseract" else None), \ mock.patch.object(ex, "extract_image_tool", boom): s, idx = self.run_stage(use_external=True) self.assertEqual(idx["scan.png"]["status"], "error") self.assertIn("no language data", idx["scan.png"]["error"]) self.assertEqual(s["counts"]["error"], 2) # + broken.docx def test_oversized_csv_field_is_an_error_entry_not_a_crash(self): # csv.Error is not a ValueError; a >128 KiB field used to escape the # handler and take the whole sweep down put(self.root, "huge.csv", "a,b\n1," + "x" * 200_000 + "\n") s, idx = self.run_stage(max_file_bytes=1_000_000) self.assertEqual(idx["huge.csv"]["status"], "error") self.assertIn("Error", idx["huge.csv"]["error"]) self.assertEqual(s["counts"]["error"], 2) # + broken.docx def test_requires_walk_first(self): with self.assertRaises(FileNotFoundError): ex.extract(manifest(self.root), self.out) def test_index_sorted_and_normalised(self): put(self.root, "crlf.txt", "one\r\ntwo\r\n\x00") _, idx = self.run_stage() self.assertEqual(self.text_of(idx["crlf.txt"]), "one\ntwo\n") paths = list(idx) self.assertEqual(paths, sorted(paths)) class RunToolTests(unittest.TestCase): def test_run_tool_ok_and_failure(self): import sys out = ex.run_tool([sys.executable, "-c", "print('hi')"]) self.assertEqual(out.strip(), b"hi") with self.assertRaises(RuntimeError) as cm: ex.run_tool([sys.executable, "-c", "import sys; sys.stderr.write('bad\\n'); sys.exit(3)"]) self.assertIn("exited 3: bad", str(cm.exception)) with self.assertRaises(RuntimeError): ex.run_tool(["/nonexistent/tool"]) if __name__ == "__main__": unittest.main()