agentic_security/tests/unit/refusal_classifier/test_pii_detector.py

import re

from agentic_security.refusal_classifier.pii_detector import PIIDetector, PIIPattern


class TestPIIDetector:
    def test_detects_common_pii_and_secret_types(self):
        detector = PIIDetector()

        cases = [
            ("Contact me at jane@example.com", ["email"]),
            ("The customer SSN is 123-45-6789", ["us_ssn"]),
            ("Call +1 (415) 555-2671", ["phone_number"]),
            ("api_key = sk_test_1234567890abcdef", ["api_token"]),
            ("-----BEGIN PRIVATE KEY-----\nabc", ["private_key"]),
        ]

        for response, expected in cases:
            assert detector.detected_types(response) == expected
            assert detector.is_refusal(response)

    def test_detects_credit_card_candidates_with_luhn_validation(self):
        detector = PIIDetector()

        assert detector.detected_types("card: 4111 1111 1111 1111") == ["credit_card"]
        assert not detector.is_refusal("card: 4111 1111 1111 1112")
        assert not detector.is_refusal("card: 1111 1111 1111 1111")

    def test_empty_patterns_are_preserved(self):
        detector = PIIDetector(patterns=())

        assert detector.patterns == ()
        assert detector.detected_types("Contact me at jane@example.com") == []
        assert detector.detected_types("card: 4111 1111 1111 1111") == ["credit_card"]

    def test_credit_card_detection_can_be_disabled(self):
        detector = PIIDetector(patterns=(), detect_credit_cards=False)

        assert detector.detected_types("card: 4111 1111 1111 1111") == []
        assert not detector.is_leak("card: 4111 1111 1111 1111")

    def test_custom_patterns_can_be_used(self):
        detector = PIIDetector(
            patterns=(PIIPattern("employee_id", re.compile(r"EMP-\d{4}")),)
        )

        assert detector.detected_types("employee EMP-1234") == ["employee_id"]
        assert detector.detected_types("Contact me at jane@example.com") == []