info-privacy/tests/test_api.py

# tests/test_api.py
import pytest
import docx
import io
import json
import openpyxl
from fastapi.testclient import TestClient
from info_privacy.api.main import app

client = TestClient(app)


def test_health():
    r = client.get("/api/v1/health")
    assert r.status_code == 200
    assert r.json()["status"] == "ok"


def test_types():
    r = client.get("/api/v1/types")
    assert r.status_code == 200
    assert "id_card" in r.json()["entity_types"]


def test_analyze_docx(tmp_path):
    doc = docx.Document()
    doc.add_paragraph("身份证：110101199001011234 电话：13812345678")
    p = tmp_path / "test.docx"
    doc.save(str(p))
    with open(p, "rb") as f:
        r = client.post("/api/v1/analyze",
                        files={"file": ("test.docx", f, "application/octet-stream")})
    assert r.status_code == 200
    data = r.json()
    assert data["blocked"] is False
    types = {e["type"] for e in data["entities"]}
    assert "id_card" in types or "phone" in types


def test_analyze_classified(tmp_path):
    doc = docx.Document()
    doc.add_paragraph("【机密】内部文件")
    p = tmp_path / "secret.docx"
    doc.save(str(p))
    with open(p, "rb") as f:
        r = client.post("/api/v1/analyze",
                        files={"file": ("secret.docx", f, "application/octet-stream")})
    assert r.status_code == 200
    assert r.json()["blocked"] is True


def test_redact_docx(tmp_path):
    doc = docx.Document()
    doc.add_paragraph("身份证：110101199001011234 地址：北京市")
    p = tmp_path / "test.docx"
    doc.save(str(p))
    cfg = json.dumps({"redact_types": ["id_card"]})
    with open(p, "rb") as f:
        r = client.post("/api/v1/redact",
                        files={"file": ("test.docx", f, "application/octet-stream")},
                        data={"config": cfg})
    assert r.status_code == 200
    result_doc = docx.Document(io.BytesIO(r.content))
    all_text = " ".join(para.text for para in result_doc.paragraphs)
    assert "110101199001011234" not in all_text


def test_redact_returns_x_security_report_header(tmp_path):
    """redact 响应头应包含 X-Security-Report，内含 JSON 格式遮罩统计。"""
    doc = docx.Document()
    doc.add_paragraph("身份证：110101199001011234")
    p = tmp_path / "test.docx"
    doc.save(str(p))
    cfg = json.dumps({"redact_types": ["id_card"]})
    with open(p, "rb") as f:
        r = client.post("/api/v1/redact",
                        files={"file": ("test.docx", f, "application/octet-stream")},
                        data={"config": cfg})
    assert r.status_code == 200
    assert "x-security-report" in r.headers
    report = json.loads(r.headers["x-security-report"])
    assert isinstance(report, dict)


def test_analyze_xlsx(tmp_path):
    """analyze 端点应支持 .xlsx 格式并返回检测结果。"""
    wb = openpyxl.Workbook()
    ws = wb.active
    ws["A1"] = "姓名"
    ws["B1"] = "身份证"
    ws["A2"] = "张三"
    ws["B2"] = "110101199001011234"
    p = tmp_path / "test.xlsx"
    wb.save(str(p))
    with open(p, "rb") as f:
        r = client.post("/api/v1/analyze",
                        files={"file": ("test.xlsx", f, "application/octet-stream")})
    assert r.status_code == 200
    data = r.json()
    assert any(e["type"] == "id_card" for e in data["entities"])


def test_analyze_no_pii_returns_normal(tmp_path):
    """不含任何 PII 的文档应返回 normal 分类且 blocked=False。"""
    doc = docx.Document()
    doc.add_paragraph("本季度销售额同比增长15%，各项指标均达标。")
    p = tmp_path / "clean.docx"
    doc.save(str(p))
    with open(p, "rb") as f:
        r = client.post("/api/v1/analyze",
                        files={"file": ("clean.docx", f, "application/octet-stream")})
    assert r.status_code == 200
    data = r.json()
    assert data["classification"] == "normal"
    assert data["blocked"] is False


def test_analyze_summary_matches_entities(tmp_path):
    """analyze 响应中 summary 各类型计数应与 entities 列表精确一致。"""
    doc = docx.Document()
    doc.add_paragraph("身份证：110101199001011234 电话：13812345678")
    p = tmp_path / "test.docx"
    doc.save(str(p))
    with open(p, "rb") as f:
        r = client.post("/api/v1/analyze",
                        files={"file": ("test.docx", f, "application/octet-stream")})
    assert r.status_code == 200
    data = r.json()
    for etype, count in data["summary"].items():
        actual = sum(1 for e in data["entities"] if e["type"] == etype)
        assert actual == count, f"{etype}: summary={count}，实际 entities={actual}"


def test_redact_classified_returns_403(tmp_path):
    """对保密文档执行 redact 应返回 403。"""
    doc = docx.Document()
    doc.add_paragraph("【机密】内部绝密资料")
    p = tmp_path / "secret.docx"
    doc.save(str(p))
    cfg = json.dumps({"redact_types": ["id_card"]})
    with open(p, "rb") as f:
        r = client.post("/api/v1/redact",
                        files={"file": ("secret.docx", f, "application/octet-stream")},
                        data={"config": cfg})
    assert r.status_code == 403