117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
"""tmkit 全链路测试。"""
|
||
|
||
import fitz
|
||
import pytest
|
||
|
||
import sys, os
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from tmkit.splitter import classify_pages
|
||
from tmkit.extractor import extract_page_words, extract_pdf_words
|
||
from tmkit.detector import PIIDetector, Rule, default_detector, _luhn_ok, _idcard_ok
|
||
from tmkit.redactor import redact_pdf
|
||
from .conftest import VALID_ID, BANK_CARD, PHONE, EMAIL
|
||
|
||
|
||
def _scan(pdf_bytes):
|
||
det = default_detector()
|
||
hits = []
|
||
with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
|
||
for pno, page in enumerate(doc):
|
||
words = extract_page_words(page)
|
||
hits += det.scan_page_words(words, pno)
|
||
return hits
|
||
|
||
|
||
# ---------- splitter ----------
|
||
|
||
def test_splitter_native_vs_scanned(pdf_plain, pdf_scanned):
|
||
assert classify_pages(pdf_plain) == ["native"]
|
||
assert classify_pages(pdf_scanned) == ["scanned"]
|
||
|
||
|
||
# ---------- extractor ----------
|
||
|
||
def test_extractor_words_with_bbox(pdf_plain):
|
||
with fitz.open(stream=pdf_plain, filetype="pdf") as doc:
|
||
words = extract_page_words(doc[0])
|
||
assert words and all("text" in w and len(w["bbox"]) == 4 for w in words)
|
||
joined = "".join(w["text"] for w in words)
|
||
assert "判决书" in joined
|
||
|
||
|
||
def test_extractor_scanned_raises(pdf_scanned):
|
||
with pytest.raises(NotImplementedError):
|
||
extract_pdf_words(pdf_scanned)
|
||
|
||
|
||
# ---------- detector 校验算法 ----------
|
||
|
||
def test_idcard_validator():
|
||
assert _idcard_ok(VALID_ID)
|
||
assert not _idcard_ok(VALID_ID[:-1] + ("0" if VALID_ID[-1] != "0" else "1"))
|
||
|
||
|
||
def test_luhn_validator():
|
||
assert _luhn_ok(BANK_CARD)
|
||
assert not _luhn_ok(BANK_CARD[:-1] + ("0" if BANK_CARD[-1] != "0" else "1"))
|
||
|
||
|
||
# ---------- detector 全链路（文本 + 坐标） ----------
|
||
|
||
def test_detector_types_found(pdf_pii):
|
||
hits = _scan(pdf_pii)
|
||
types = {h["type"] for h in hits}
|
||
assert types == {"id_card", "phone", "bank_card", "email"}
|
||
texts = " ".join(h["text"] for h in hits)
|
||
assert VALID_ID in texts and PHONE in texts and BANK_CARD in texts and EMAIL in texts
|
||
for h in hits:
|
||
assert set(h) == {"type", "text", "page", "bbox"}
|
||
assert h["page"] == 0 and h["bbox"][2] > h["bbox"][0]
|
||
|
||
|
||
def test_detector_no_false_positive_on_plain(pdf_plain):
|
||
assert _scan(pdf_plain) == []
|
||
|
||
|
||
def test_detector_custom_rule():
|
||
det = PIIDetector(rules=[])
|
||
det.register(Rule("order_no", __import__("re").compile(r"（2026）京01民初1号")))
|
||
hits = det.scan_text("案号：（2026）京01民初1号。")
|
||
assert hits == [{"type": "order_no", "text": "（2026）京01民初1号"}]
|
||
|
||
|
||
# ---------- redactor 全链路 ----------
|
||
|
||
@pytest.mark.parametrize("method", ["redact", "draw"])
|
||
def test_redact_pipeline(pdf_pii, method):
|
||
hits = _scan(pdf_pii)
|
||
out = redact_pdf(pdf_pii, hits, method=method)
|
||
with fitz.open(stream=out, filetype="pdf") as doc:
|
||
assert doc.page_count == 1 # 页数不变
|
||
text = doc[0].get_text()
|
||
if method == "redact":
|
||
# PII 彻底删除
|
||
assert VALID_ID not in text and PHONE not in text
|
||
assert BANK_CARD not in text and EMAIL not in text
|
||
# 非敏感文本保留
|
||
assert "张三" in text and "脱敏" in text
|
||
else:
|
||
# draw 模式只画黑条，文本层仍在（已知限制，测试如实断言）
|
||
assert VALID_ID in text
|
||
|
||
|
||
def test_redact_idempotent_on_clean_pdf(pdf_plain):
|
||
out = redact_pdf(pdf_plain, [])
|
||
with fitz.open(stream=out, filetype="pdf") as doc:
|
||
assert "判决书" in doc[0].get_text()
|
||
|
||
|
||
def test_redact_black_box_visible(pdf_pii):
|
||
hits = _scan(pdf_pii)
|
||
out = redact_pdf(pdf_pii, hits, method="redact")
|
||
# apply_redactions(fill) 后区域为黑：检查像素
|
||
with fitz.open(stream=out, filetype="pdf") as doc:
|
||
pix = doc[0].get_pixmap(clip=fitz.Rect(*hits[0]["bbox"]))
|
||
assert pix.pixel(2, 2) == (0, 0, 0)
|