info-privacy/tests/test_detectors.py

import pytest
from info_privacy.models import TextBlock, EntityType
from info_privacy.detectors.regex_detector import RegexDetector


@pytest.fixture
def detector():
    return RegexDetector("configs/pii_rules.yaml")


def test_detect_id_card(detector):
    block = TextBlock(text="姓名：张三 身份证：110101199001011234", bbox=[0, 0, 100, 20], page=1, layer="text")
    entities = detector.detect([block])
    types = [e.type for e in entities]
    assert EntityType.ID_CARD in types


def test_detect_phone(detector):
    block = TextBlock(text="联系电话：13812345678", bbox=[0, 0, 100, 20], page=1, layer="text")
    entities = detector.detect([block])
    types = [e.type for e in entities]
    assert EntityType.PHONE in types


def test_detect_email(detector):
    block = TextBlock(text="邮箱：test@example.com", bbox=[0, 0, 100, 20], page=1, layer="text")
    entities = detector.detect([block])
    assert any(e.type == EntityType.EMAIL for e in entities)


def test_no_false_positive_short_number(detector):
    block = TextBlock(text="编号：12345", bbox=[0, 0, 100, 20], page=1, layer="text")
    entities = detector.detect([block])
    assert not any(e.type == EntityType.ID_CARD for e in entities)


def test_entity_bbox_matches_block(detector):
    block = TextBlock(text="电话：13812345678", bbox=[10, 20, 200, 40], page=2, layer="text")
    entities = detector.detect([block])
    assert entities[0].page == 2
    assert entities[0].bbox == [10, 20, 200, 40]


# NER 检测器测试
from info_privacy.detectors.ner_detector import NERDetector

@pytest.fixture
def ner():
    return NERDetector("configs/pii_rules.yaml")

def test_detect_name(ner):
    block = TextBlock(text="申请人：张伟，联系地址如下", bbox=[0,0,200,20], page=1, layer="text")
    entities = ner.detect([block])
    assert any(e.type == EntityType.NAME for e in entities)

def test_detect_address(ner):
    block = TextBlock(text="住址：北京市朝阳区建国路88号", bbox=[0,0,300,20], page=1, layer="text")
    entities = ner.detect([block])
    assert any(e.type == EntityType.ADDRESS for e in entities)

def test_no_name_without_surname(ner):
    block = TextBlock(text="操作系统版本3.1", bbox=[0,0,200,20], page=1, layer="text")
    entities = ner.detect([block])
    assert not any(e.type == EntityType.NAME for e in entities)


def test_ner_no_address_false_positive(ner):
    """普通正文不含地址触发词，不应检测到 address。"""
    block = TextBlock(text="本季度产品销售额增长20%", bbox=[0,0,300,20], page=1, layer="text")
    entities = ner.detect([block])
    assert not any(e.type == EntityType.ADDRESS for e in entities)


def test_ner_name_and_address_same_block(ner):
    """同一段文本中姓名和地址应同时被检测。"""
    block = TextBlock(
        text="申请人：王芳，地址：上海市浦东新区张江路100号",
        bbox=[0, 0, 400, 20], page=1, layer="text",
    )
    entities = ner.detect([block])
    types = {e.type for e in entities}
    assert EntityType.NAME in types
    assert EntityType.ADDRESS in types


# RegexDetector 新增覆盖

def test_detect_bank_card_16_digit(detector):
    """16 位纯银行卡号（非身份证格式）应被识别为 bank_card。"""
    block = TextBlock(text="卡号：6222021302001002", bbox=[0,0,200,20], page=1, layer="text")
    entities = detector.detect([block])
    types = [e.type for e in entities]
    assert EntityType.BANK_CARD in types
    assert EntityType.ID_CARD not in types


def test_id_card_not_double_counted_as_bank_card(detector):
    """18 位身份证号不应同时被识别为 bank_card（去重逻辑）。"""
    block = TextBlock(text="证件号：110101199001011234", bbox=[0,0,200,20], page=1, layer="text")
    entities = detector.detect([block])
    id_entities = [e for e in entities if e.type == EntityType.ID_CARD]
    bank_overlap = [e for e in entities
                    if e.type == EntityType.BANK_CARD and e.value == "110101199001011234"]
    assert len(id_entities) == 1
    assert len(bank_overlap) == 0


def test_detect_id_card_ending_in_x(detector):
    """末位为 X 的身份证号应被正确识别，且 value 保留末位 X。"""
    block = TextBlock(text="证件：11010119900101110X", bbox=[0,0,200,20], page=1, layer="text")
    entities = detector.detect([block])
    id_entities = [e for e in entities if e.type == EntityType.ID_CARD]
    assert len(id_entities) == 1
    assert id_entities[0].value.endswith("X")


def test_detect_license_plate(detector):
    """有效的中国车牌应被识别为 license_plate。"""
    block = TextBlock(text="车辆：京A12345", bbox=[0,0,200,20], page=1, layer="text")
    entities = detector.detect([block])
    assert any(e.type == EntityType.LICENSE_PLATE for e in entities)


def test_entity_security_level_high(detector):
    """高风险类型 id_card 的 security_level 应为 high。"""
    block = TextBlock(text="110101199001011234", bbox=[0,0,200,20], page=1, layer="text")
    entities = detector.detect([block])
    id_e = next(e for e in entities if e.type == EntityType.ID_CARD)
    assert id_e.security_level == "high"


def test_entity_security_level_medium(detector):
    """中风险类型 email 的 security_level 应为 medium。"""
    block = TextBlock(text="user@example.com", bbox=[0,0,200,20], page=1, layer="text")
    entities = detector.detect([block])
    email_e = next(e for e in entities if e.type == EntityType.EMAIL)
    assert email_e.security_level == "medium"


def test_multiple_entities_same_block(detector):
    """同一文本块中同时含手机号和邮箱，两者均应被检测。"""
    block = TextBlock(
        text="联系方式：13812345678 / admin@corp.com",
        bbox=[0, 0, 400, 20], page=1, layer="text",
    )
    entities = detector.detect([block])
    types = {e.type for e in entities}
    assert EntityType.PHONE in types
    assert EntityType.EMAIL in types


# 人脸检测器测试
import os
import numpy as np
import cv2

FACE_MODEL = "/data/rockchip/mediapipe/models/rknn/face_detection_short_range_rk3588.rknn"

def _rknn_available():
    try:
        import rknn
        return True
    except ImportError:
        try:
            import rknnlite
            return True
        except ImportError:
            return False

def test_face_detector_skip_if_no_model():
    if not os.path.exists(FACE_MODEL) or not _rknn_available():
        pytest.skip("人脸RKNN模型不存在或 RKNN 运行时不可用")
    from info_privacy.detectors.face_detector import FaceDetector
    detector = FaceDetector(FACE_MODEL)
    assert detector is not None

def test_face_detector_returns_entities():
    if not os.path.exists(FACE_MODEL) or not _rknn_available():
        pytest.skip("人脸RKNN模型不存在或 RKNN 运行时不可用")
    from info_privacy.detectors.face_detector import FaceDetector
    detector = FaceDetector(FACE_MODEL)
    # 白色空图像，无人脸
    img = np.ones((300, 300, 3), dtype=np.uint8) * 255
    entities = detector.detect_in_image(img, page=1)
    assert isinstance(entities, list)