418 lines
17 KiB
Python
418 lines
17 KiB
Python
|
"""Unit tests for the documents-ingest PDF extractor — no DB, no external services."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from email import encoders
|
||
|
|
from email.mime.application import MIMEApplication
|
||
|
|
from email.mime.multipart import MIMEMultipart
|
||
|
|
from email.mime.text import MIMEText
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from documents_ingest.extractor import (
|
||
|
|
Candidate,
|
||
|
|
build_consume_name,
|
||
|
|
find_pdf_parts,
|
||
|
|
load_registry,
|
||
|
|
pdf_candidates_from_entities,
|
||
|
|
process_candidate,
|
||
|
|
run,
|
||
|
|
sanitize_filename,
|
||
|
|
save_registry,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _build_eml(parts: list[tuple[str, str, bytes]]) -> bytes:
|
||
|
|
"""parts: list of (filename, content_type, bytes). Builds a multipart .eml."""
|
||
|
|
msg = MIMEMultipart()
|
||
|
|
msg["From"] = "alice@example.com"
|
||
|
|
msg["To"] = "bob@example.com"
|
||
|
|
msg["Subject"] = "Invoice attached"
|
||
|
|
msg["Message-ID"] = "<test@example.com>"
|
||
|
|
msg["Date"] = "Tue, 10 Jun 2025 12:00:00 +0000"
|
||
|
|
msg.attach(MIMEText("See attached", "plain"))
|
||
|
|
for filename, content_type, data in parts:
|
||
|
|
maintype, subtype = content_type.split("/", 1)
|
||
|
|
att = MIMEApplication(data, _subtype=subtype) if maintype == "application" else None
|
||
|
|
if att is None:
|
||
|
|
from email.mime.base import MIMEBase
|
||
|
|
att = MIMEBase(maintype, subtype)
|
||
|
|
att.set_payload(data)
|
||
|
|
encoders.encode_base64(att)
|
||
|
|
att.add_header("Content-Disposition", "attachment", filename=filename)
|
||
|
|
msg.attach(att)
|
||
|
|
return msg.as_bytes()
|
||
|
|
|
||
|
|
|
||
|
|
PDF_BYTES = b"%PDF-1.4 fake pdf content for testing"
|
||
|
|
PDF_SHA256 = hashlib.sha256(PDF_BYTES).hexdigest()
|
||
|
|
|
||
|
|
|
||
|
|
class TestSanitizeFilename:
|
||
|
|
def test_strips_path_components(self):
|
||
|
|
assert sanitize_filename("../../etc/passwd.pdf") == "passwd.pdf"
|
||
|
|
|
||
|
|
def test_replaces_unsafe_chars(self):
|
||
|
|
assert sanitize_filename("faktura #123 (final).pdf") == "faktura_123_final_.pdf"
|
||
|
|
|
||
|
|
def test_appends_pdf_extension_if_missing(self):
|
||
|
|
assert sanitize_filename("report").endswith(".pdf")
|
||
|
|
|
||
|
|
def test_empty_name_falls_back(self):
|
||
|
|
assert sanitize_filename("") == "attachment.pdf"
|
||
|
|
|
||
|
|
def test_only_unsafe_chars_falls_back(self):
|
||
|
|
assert sanitize_filename("///...") == "attachment.pdf"
|
||
|
|
|
||
|
|
def test_truncates_long_names(self):
|
||
|
|
long_name = ("a" * 300) + ".pdf"
|
||
|
|
result = sanitize_filename(long_name)
|
||
|
|
assert len(result) <= 150
|
||
|
|
assert result.endswith(".pdf")
|
||
|
|
|
||
|
|
def test_preserves_reasonable_name(self):
|
||
|
|
assert sanitize_filename("faktura_2026-06.pdf") == "faktura_2026-06.pdf"
|
||
|
|
|
||
|
|
|
||
|
|
class TestBuildConsumeName:
|
||
|
|
def test_basic_name(self):
|
||
|
|
ts = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||
|
|
used: set[str] = set()
|
||
|
|
name = build_consume_name(ts, "invoice.pdf", "abc123", used)
|
||
|
|
assert name == "2026-06-01_invoice.pdf"
|
||
|
|
|
||
|
|
def test_collision_gets_hash_suffix(self):
|
||
|
|
ts = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||
|
|
used: set[str] = set()
|
||
|
|
first = build_consume_name(ts, "invoice.pdf", "aaaa1111", used)
|
||
|
|
second = build_consume_name(ts, "invoice.pdf", "bbbb2222", used)
|
||
|
|
assert first != second
|
||
|
|
assert second == "2026-06-01_invoice_bbbb2222.pdf"
|
||
|
|
|
||
|
|
def test_names_are_registered_in_used_set(self):
|
||
|
|
ts = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||
|
|
used: set[str] = set()
|
||
|
|
name = build_consume_name(ts, "invoice.pdf", "abc123", used)
|
||
|
|
assert name in used
|
||
|
|
|
||
|
|
|
||
|
|
class TestPdfCandidatesFromEntities:
|
||
|
|
ts = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||
|
|
|
||
|
|
def test_filters_by_content_type_and_size(self):
|
||
|
|
entities = [
|
||
|
|
{"type": "attachment", "content_type": "application/pdf", "size": 60000,
|
||
|
|
"filename": "big.pdf", "sha256": "x"},
|
||
|
|
{"type": "attachment", "content_type": "application/pdf", "size": 1000,
|
||
|
|
"filename": "small.pdf", "sha256": "y"},
|
||
|
|
{"type": "attachment", "content_type": "image/png", "size": 100000,
|
||
|
|
"filename": "photo.png", "sha256": "z"},
|
||
|
|
]
|
||
|
|
candidates = pdf_candidates_from_entities("env1", self.ts, entities, min_size=50000)
|
||
|
|
assert len(candidates) == 1
|
||
|
|
assert candidates[0].filename == "big.pdf"
|
||
|
|
|
||
|
|
def test_skips_entries_without_sha256(self):
|
||
|
|
entities = [
|
||
|
|
{"type": "attachment", "content_type": "application/pdf", "size": 60000,
|
||
|
|
"filename": "big.pdf"},
|
||
|
|
]
|
||
|
|
assert pdf_candidates_from_entities("env1", self.ts, entities, min_size=50000) == []
|
||
|
|
|
||
|
|
def test_ignores_non_attachment_entries(self):
|
||
|
|
entities = [{"type": "inline", "content_type": "application/pdf", "size": 60000}]
|
||
|
|
assert pdf_candidates_from_entities("env1", self.ts, entities, min_size=50000) == []
|
||
|
|
|
||
|
|
|
||
|
|
class TestFindPdfParts:
|
||
|
|
def test_finds_pdf_part(self):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
parts = find_pdf_parts(raw)
|
||
|
|
assert parts == [("report.pdf", PDF_BYTES)]
|
||
|
|
|
||
|
|
def test_no_pdf_parts_returns_empty(self):
|
||
|
|
raw = _build_eml([("report.png", "image/png", b"not a pdf")])
|
||
|
|
assert find_pdf_parts(raw) == []
|
||
|
|
|
||
|
|
def test_multiple_pdf_attachments_all_returned(self):
|
||
|
|
raw = _build_eml([
|
||
|
|
("a.pdf", "application/pdf", PDF_BYTES),
|
||
|
|
("b.pdf", "application/pdf", b"other pdf"),
|
||
|
|
("photo.png", "image/png", b"png bytes"),
|
||
|
|
])
|
||
|
|
parts = find_pdf_parts(raw)
|
||
|
|
assert ("a.pdf", PDF_BYTES) in parts
|
||
|
|
assert ("b.pdf", b"other pdf") in parts
|
||
|
|
assert len(parts) == 2
|
||
|
|
|
||
|
|
def test_decodes_rfc2047_encoded_filenames(self):
|
||
|
|
# Regression: the manifest (built by a different parser at import time)
|
||
|
|
# stores the RAW encoded-word filename; email.policy.default decodes it
|
||
|
|
# to real unicode here. process_candidate() must not require these to
|
||
|
|
# match byte-for-byte — sha256 is the authority (see TestProcessCandidate).
|
||
|
|
raw = _build_eml([
|
||
|
|
("=?UTF-8?b?ZmFrdHVyYV/FunJvZGxvLnBkZg==?=", "application/pdf", PDF_BYTES),
|
||
|
|
])
|
||
|
|
parts = find_pdf_parts(raw)
|
||
|
|
assert len(parts) == 1
|
||
|
|
filename, payload = parts[0]
|
||
|
|
assert filename == "faktura_źrodlo.pdf"
|
||
|
|
assert payload == PDF_BYTES
|
||
|
|
|
||
|
|
|
||
|
|
class TestProcessCandidate:
|
||
|
|
ts = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||
|
|
|
||
|
|
def _candidate(self, sha256=PDF_SHA256, filename="report.pdf"):
|
||
|
|
return Candidate(envelope_id="env1", ts=self.ts, filename=filename,
|
||
|
|
size=len(PDF_BYTES), sha256=sha256)
|
||
|
|
|
||
|
|
def test_extracted_on_sha_match(self):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
status, example = process_candidate(raw, self._candidate(), {}, set(), apply=False)
|
||
|
|
assert status == "extracted"
|
||
|
|
assert example["target_name"] == "2026-06-01_report.pdf"
|
||
|
|
assert example["sha256"] == PDF_SHA256
|
||
|
|
|
||
|
|
def test_duplicate_when_sha_in_registry(self):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
registry = {PDF_SHA256: {"envelope_id": "env0"}}
|
||
|
|
status, example = process_candidate(raw, self._candidate(), registry, set(), apply=False)
|
||
|
|
assert status == "duplicate"
|
||
|
|
assert example is None
|
||
|
|
|
||
|
|
def test_sha_mismatch_reported_and_skipped(self):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
candidate = self._candidate(sha256="0" * 64)
|
||
|
|
status, example = process_candidate(raw, candidate, {}, set(), apply=False)
|
||
|
|
assert status == "sha_mismatch"
|
||
|
|
assert example is None
|
||
|
|
|
||
|
|
def test_parse_error_when_attachment_missing_from_mime(self):
|
||
|
|
raw = _build_eml([("other.pdf", "application/pdf", b"unrelated pdf bytes")])
|
||
|
|
status, example = process_candidate(raw, self._candidate(), {}, set(), apply=False)
|
||
|
|
assert status == "parse_error"
|
||
|
|
assert example is None
|
||
|
|
|
||
|
|
def test_sha_match_wins_over_filename_mismatch(self):
|
||
|
|
# Regression: manifest filename is the raw RFC 2047 encoded-word form,
|
||
|
|
# but the freshly re-parsed MIME part decodes to real unicode text.
|
||
|
|
# sha256 is the proof of identity — this must still extract.
|
||
|
|
raw = _build_eml([
|
||
|
|
("=?UTF-8?b?ZmFrdHVyYV/FunJvZGxvLnBkZg==?=", "application/pdf", PDF_BYTES),
|
||
|
|
])
|
||
|
|
candidate = self._candidate(filename="=?UTF-8?b?ZmFrdHVyYV/FunJvZGxvLnBkZg==?=")
|
||
|
|
status, example = process_candidate(raw, candidate, {}, set(), apply=False)
|
||
|
|
assert status == "extracted"
|
||
|
|
assert example["sha256"] == PDF_SHA256
|
||
|
|
# Consume/ filename is built from the decoded name, not the raw encoded-word.
|
||
|
|
assert example["filename"] == "faktura_źrodlo.pdf"
|
||
|
|
assert "=?UTF-8?" not in example["target_name"]
|
||
|
|
|
||
|
|
def test_apply_mode_includes_payload_for_writing(self):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
status, example = process_candidate(raw, self._candidate(), {}, set(), apply=True)
|
||
|
|
assert status == "extracted"
|
||
|
|
assert example["payload"] == PDF_BYTES
|
||
|
|
|
||
|
|
def test_dry_run_mode_has_no_payload(self):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
status, example = process_candidate(raw, self._candidate(), {}, set(), apply=False)
|
||
|
|
assert example["payload"] is None
|
||
|
|
|
||
|
|
|
||
|
|
class TestRegistry:
|
||
|
|
def test_missing_file_returns_empty_dict(self, tmp_path):
|
||
|
|
assert load_registry(tmp_path / "nope.json") == {}
|
||
|
|
|
||
|
|
def test_round_trip(self, tmp_path):
|
||
|
|
path = tmp_path / "sub" / "registry.json"
|
||
|
|
registry = {"abc123": {"envelope_id": "env1", "consume_name": "x.pdf"}}
|
||
|
|
save_registry(path, registry)
|
||
|
|
assert load_registry(path) == registry
|
||
|
|
|
||
|
|
def test_save_creates_parent_dirs(self, tmp_path):
|
||
|
|
path = tmp_path / "a" / "b" / "registry.json"
|
||
|
|
save_registry(path, {})
|
||
|
|
assert path.is_file()
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeConn:
|
||
|
|
def __init__(self, rows):
|
||
|
|
self._rows = rows
|
||
|
|
|
||
|
|
async def fetch(self, query, *params):
|
||
|
|
return self._rows
|
||
|
|
|
||
|
|
async def close(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def _row(envelope_id, ts, raw_ref, entities):
|
||
|
|
return {"id": envelope_id, "raw_ref": raw_ref, "ts": ts, "entities": json.dumps(entities)}
|
||
|
|
|
||
|
|
|
||
|
|
def _pdf_entity(filename="report.pdf", size=60000, sha256=PDF_SHA256):
|
||
|
|
return {"type": "attachment", "content_type": "application/pdf", "size": size,
|
||
|
|
"filename": filename, "sha256": sha256}
|
||
|
|
|
||
|
|
|
||
|
|
class TestRun:
|
||
|
|
ts = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||
|
|
|
||
|
|
def _setup_archive(self, tmp_path, raw_ref, raw_bytes):
|
||
|
|
archive_root = tmp_path / "archive"
|
||
|
|
eml_path = archive_root / raw_ref
|
||
|
|
eml_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
eml_path.write_bytes(raw_bytes)
|
||
|
|
return archive_root
|
||
|
|
|
||
|
|
def _patch_connect(self, monkeypatch, rows):
|
||
|
|
async def _fake_connect(dsn):
|
||
|
|
return _FakeConn(rows)
|
||
|
|
monkeypatch.setattr("documents_ingest.extractor.asyncpg.connect", _fake_connect)
|
||
|
|
|
||
|
|
async def test_dry_run_does_not_write_files_or_registry(self, tmp_path, monkeypatch):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
archive_root = self._setup_archive(tmp_path, "gmail/2026/06/m1.eml", raw)
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/m1.eml", [_pdf_entity()])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
stats, examples = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats["extracted"] == 1
|
||
|
|
assert len(examples) == 1
|
||
|
|
assert list(consume_dir.iterdir()) == []
|
||
|
|
assert not registry_path.exists()
|
||
|
|
|
||
|
|
async def test_apply_writes_file_and_registry(self, tmp_path, monkeypatch):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
archive_root = self._setup_archive(tmp_path, "gmail/2026/06/m1.eml", raw)
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/m1.eml", [_pdf_entity()])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
stats, examples = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats["extracted"] == 1
|
||
|
|
written = list(consume_dir.iterdir())
|
||
|
|
assert len(written) == 1
|
||
|
|
assert written[0].read_bytes() == PDF_BYTES
|
||
|
|
registry = load_registry(registry_path)
|
||
|
|
assert PDF_SHA256 in registry
|
||
|
|
|
||
|
|
async def test_idempotent_second_apply_run_skips_duplicate(self, tmp_path, monkeypatch):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
archive_root = self._setup_archive(tmp_path, "gmail/2026/06/m1.eml", raw)
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/m1.eml", [_pdf_entity()])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
await run(dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=True)
|
||
|
|
stats2, examples2 = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats2["extracted"] == 0
|
||
|
|
assert stats2["skipped_duplicate"] == 1
|
||
|
|
assert len(list(consume_dir.iterdir())) == 1
|
||
|
|
|
||
|
|
async def test_multiple_candidates_in_one_envelope(self, tmp_path, monkeypatch):
|
||
|
|
pdf2 = b"%PDF-1.4 second document"
|
||
|
|
sha2 = hashlib.sha256(pdf2).hexdigest()
|
||
|
|
raw = _build_eml([
|
||
|
|
("a.pdf", "application/pdf", PDF_BYTES),
|
||
|
|
("b.pdf", "application/pdf", pdf2),
|
||
|
|
])
|
||
|
|
archive_root = self._setup_archive(tmp_path, "gmail/2026/06/m1.eml", raw)
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/m1.eml", [
|
||
|
|
_pdf_entity(filename="a.pdf", sha256=PDF_SHA256),
|
||
|
|
_pdf_entity(filename="b.pdf", sha256=sha2),
|
||
|
|
])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
stats, examples = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats["extracted"] == 2
|
||
|
|
assert len(list(consume_dir.iterdir())) == 2
|
||
|
|
|
||
|
|
async def test_sha_mismatch_is_not_written(self, tmp_path, monkeypatch):
|
||
|
|
raw = _build_eml([("report.pdf", "application/pdf", PDF_BYTES)])
|
||
|
|
archive_root = self._setup_archive(tmp_path, "gmail/2026/06/m1.eml", raw)
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/m1.eml", [
|
||
|
|
_pdf_entity(sha256="0" * 64),
|
||
|
|
])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
stats, examples = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats["extracted"] == 0
|
||
|
|
assert stats["skipped_sha_mismatch"] == 1
|
||
|
|
assert list(consume_dir.iterdir()) == []
|
||
|
|
|
||
|
|
async def test_missing_eml_file_counts_as_error(self, tmp_path, monkeypatch):
|
||
|
|
archive_root = tmp_path / "archive"
|
||
|
|
archive_root.mkdir()
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/missing.eml", [_pdf_entity()])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
stats, examples = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats["errors"] == 1
|
||
|
|
assert stats["extracted"] == 0
|
||
|
|
|
||
|
|
async def test_envelope_without_qualifying_entities_is_scanned_but_skipped(self, tmp_path, monkeypatch):
|
||
|
|
archive_root = tmp_path / "archive"
|
||
|
|
archive_root.mkdir()
|
||
|
|
consume_dir = tmp_path / "consume"
|
||
|
|
consume_dir.mkdir()
|
||
|
|
registry_path = tmp_path / "registry.json"
|
||
|
|
rows = [_row("env1", self.ts, "gmail/2026/06/m1.eml", [
|
||
|
|
{"type": "attachment", "content_type": "image/png", "size": 100, "filename": "x.png", "sha256": "z"},
|
||
|
|
])]
|
||
|
|
self._patch_connect(monkeypatch, rows)
|
||
|
|
|
||
|
|
stats, examples = await run(
|
||
|
|
dsn="postgresql://fake", archive_root=archive_root, consume_dir=consume_dir,
|
||
|
|
registry_path=registry_path, apply=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert stats["envelopes_scanned"] == 1
|
||
|
|
assert stats["pdf_candidates"] == 0
|
||
|
|
assert examples == []
|