qiurui
cbfe4a23dc
- Pipeline: analyze(), redact(), analyze_image() methods - API: /analyze, /redact, /analyze/frame, /analyze/frame/base64 endpoints - Detectors: regex, NER, face (RKNN NPU) - Privacy frame route added for KVM-Privacy Hub integration
90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
# tests/test_pipeline.py
|
||
import pytest
|
||
import docx
|
||
from info_privacy.pipeline import PrivacyPipeline
|
||
|
||
@pytest.fixture
|
||
def pipeline():
|
||
return PrivacyPipeline()
|
||
|
||
@pytest.fixture
|
||
def sample_docx_path(tmp_path):
|
||
doc = docx.Document()
|
||
doc.add_paragraph("姓名：张伟 身份证：110101199001011234 电话：13812345678")
|
||
p = tmp_path / "test.docx"
|
||
doc.save(str(p))
|
||
return str(p)
|
||
|
||
def test_analyze_docx(pipeline, sample_docx_path):
|
||
report = pipeline.analyze(sample_docx_path)
|
||
assert report.blocked is False
|
||
assert len(report.entities) > 0
|
||
types = {e.type.value for e in report.entities}
|
||
assert "id_card" in types or "phone" in types
|
||
|
||
def test_analyze_classified_doc(pipeline, tmp_path):
|
||
doc = docx.Document()
|
||
doc.add_paragraph("【机密】本文件仅供内部使用")
|
||
p = tmp_path / "secret.docx"
|
||
doc.save(str(p))
|
||
report = pipeline.analyze(str(p))
|
||
assert report.blocked is True
|
||
assert report.classification.value == "classified"
|
||
|
||
def test_redact_docx_removes_entity(pipeline, sample_docx_path, tmp_path):
|
||
out = str(tmp_path / "out.docx")
|
||
pipeline.redact(sample_docx_path, redact_types=["id_card"], out_path=out)
|
||
doc = docx.Document(out)
|
||
all_text = " ".join(p.text for p in doc.paragraphs)
|
||
assert "110101199001011234" not in all_text
|
||
assert "张伟" in all_text # 未指定遮罩 name，保留
|
||
|
||
|
||
def test_analyze_returns_summary_dict(pipeline, sample_docx_path):
|
||
"""DetectionReport.summary 各类型计数应与 entities 列表精确对应。"""
|
||
report = pipeline.analyze(sample_docx_path)
|
||
for etype, count in report.summary.items():
|
||
actual = sum(1 for e in report.entities if e.type.value == etype)
|
||
assert actual == count, f"{etype}: summary={count}，entities={actual}"
|
||
|
||
|
||
def test_analyze_no_pii_empty_entities(pipeline, tmp_path):
|
||
"""无 PII 的文档：entities 为空，分类为 normal，无 warning。"""
|
||
doc = docx.Document()
|
||
doc.add_paragraph("季度汇报：各项指标均达标。")
|
||
p = tmp_path / "clean.docx"
|
||
doc.save(str(p))
|
||
report = pipeline.analyze(str(p))
|
||
assert len(report.entities) == 0
|
||
assert report.classification.value == "normal"
|
||
assert report.warning is None
|
||
|
||
|
||
def test_analyze_many_entities_triggers_warning(pipeline, tmp_path):
|
||
"""5 个以上高风险实体应在 DetectionReport 中生成 warning。"""
|
||
doc = docx.Document()
|
||
ids = [
|
||
"110101199001011234", "110101199001011235", "110101199001011236",
|
||
"110101199001011237", "110101199001011238",
|
||
]
|
||
doc.add_paragraph(" ".join(ids))
|
||
p = tmp_path / "many.docx"
|
||
doc.save(str(p))
|
||
report = pipeline.analyze(str(p))
|
||
assert report.warning is not None
|
||
|
||
|
||
def test_redact_multi_type_all_removed(pipeline, tmp_path):
|
||
"""同时遮罩 id_card/phone/email 三种类型，各类值均应从文档消失。"""
|
||
doc = docx.Document()
|
||
doc.add_paragraph("身份证：110101199001011234 电话：13812345678 邮箱：user@corp.com")
|
||
p = tmp_path / "multi.docx"
|
||
doc.save(str(p))
|
||
out = str(tmp_path / "out.docx")
|
||
pipeline.redact(str(p), redact_types=["id_card", "phone", "email"], out_path=out)
|
||
result = docx.Document(out)
|
||
text = " ".join(para.text for para in result.paragraphs)
|
||
assert "110101199001011234" not in text
|
||
assert "13812345678" not in text
|
||
assert "user@corp.com" not in text
|