688 lines
34 KiB
Python
688 lines
34 KiB
Python
"""Unit + wire tests for the PII sterilization service.
|
|
|
|
Everything here is model-free: the LLM pass is either a fake callable
|
|
(pipeline tests) or a stub llama.cpp speaking the real chat-completions
|
|
wire format on loopback (service tests). Live recall quality against a
|
|
real model is verify-live.sh's job, against a running service.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
import urllib.error
|
|
import urllib.request
|
|
from unittest import mock
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
from . import detect, pipeline, service
|
|
from .vault import Vault
|
|
|
|
|
|
class DetectTests(unittest.TestCase):
|
|
def find(self, text):
|
|
return {(f.value, f.type) for f in detect.detect(text)}
|
|
|
|
def test_email(self):
|
|
self.assertIn(("sam.r@example.co.uk", "EMAIL"),
|
|
self.find("mail sam.r@example.co.uk today"))
|
|
|
|
def test_phone_forms(self):
|
|
for s in ["+1 415 555 0173", "(415) 555-0173", "415-555-0173", "415.555.0173"]:
|
|
self.assertIn((s, "PHONE"), self.find(f"call {s} now"), s)
|
|
|
|
def test_date_is_not_phone(self):
|
|
self.assertEqual(self.find("released 2026-08-03 at noon"), set())
|
|
|
|
def test_ssn_valid_and_invalid_area(self):
|
|
self.assertIn(("123-45-6789", "SSN"), self.find("ssn 123-45-6789"))
|
|
self.assertEqual(self.find("ssn 666-45-6789"), set())
|
|
self.assertEqual(self.find("ssn 000-45-6789"), set())
|
|
|
|
def test_credit_card_luhn(self):
|
|
self.assertIn(("4111 1111 1111 1111", "CREDIT_CARD"),
|
|
self.find("card 4111 1111 1111 1111 exp 12/28"))
|
|
# same shape, fails Luhn
|
|
self.assertEqual(self.find("ref 4111 1111 1111 1112"), set())
|
|
|
|
def test_ipv4_owns_dotted_span(self):
|
|
found = self.find("host at 192.168.68.54 responded")
|
|
self.assertIn(("192.168.68.54", "IP_ADDRESS"), found)
|
|
self.assertNotIn(("192.168.68.54", "PHONE"), found)
|
|
|
|
def test_ipv4_invalid_octet(self):
|
|
self.assertNotIn(("999.168.68.541", "IP_ADDRESS"), self.find("x 999.168.68.541"))
|
|
|
|
def test_ipv6_and_mac_disambiguation(self):
|
|
found = self.find("fe80::1ff:fe23:4567:890a from aa:bb:cc:dd:ee:ff")
|
|
self.assertIn(("fe80::1ff:fe23:4567:890a", "IP_ADDRESS"), found)
|
|
self.assertIn(("aa:bb:cc:dd:ee:ff", "MAC_ADDRESS"), found)
|
|
|
|
def test_sentence_ending_full_stop(self):
|
|
self.assertIn(("415-555-0173", "PHONE"), self.find("Call me at 415-555-0173."))
|
|
self.assertIn(("415.555.0173", "PHONE"), self.find("Call me at 415.555.0173."))
|
|
self.assertIn(("10.1.2.3", "IP_ADDRESS"), self.find("the gateway was 10.1.2.3."))
|
|
self.assertIn(("aa:bb:cc:dd:ee:ff", "MAC_ADDRESS"), self.find("from aa:bb:cc:dd:ee:ff."))
|
|
|
|
def test_full_stop_inside_a_longer_number_still_refuses(self):
|
|
self.assertEqual(self.find("version 10.1.2.3.4 shipped"), set())
|
|
self.assertNotIn("PHONE", {t for _, t in self.find("ref 415-555-0173.5 end")})
|
|
self.assertNotIn("PHONE", {t for _, t in self.find("file 415-555-0173.pdf")})
|
|
|
|
def test_full_stop_rule_is_the_same_for_ip_phone_and_mac(self):
|
|
# a full stop followed by a letter or digit is part of something longer
|
|
self.assertEqual(self.find("open 10.1.2.3.txt now"), set())
|
|
self.assertIn(("10.1.2.3", "IP_ADDRESS"), self.find("the host was 10.1.2.3."))
|
|
self.assertEqual(self.find("saved aa:bb:cc:dd:ee:ff.bak"), set())
|
|
self.assertEqual(self.find("saved aa:bb:cc:dd:ee:ff.txt"), set())
|
|
self.assertIn(("aa:bb:cc:dd:ee:ff", "MAC_ADDRESS"), self.find("it was aa:bb:cc:dd:ee:ff."))
|
|
|
|
# ── v1.2 checksum-anchored recognizers ──────────────────────────────
|
|
def test_iban_mod97(self):
|
|
for iban in ("GB82 WEST 1234 5698 7654 32", "DE89370400440532013000"):
|
|
self.assertIn((iban, "BANK_ACCOUNT"), self.find("pay to %s thanks" % iban), iban)
|
|
for bad in ("GB82 WEST 1234 5698 7654 33", # bad check digits
|
|
"ZZ82 WEST 1234 5698 7654 32"): # unknown country
|
|
self.assertNotIn("BANK_ACCOUNT", {t for _, t in self.find("pay to " + bad)}, bad)
|
|
|
|
def test_iban_not_reclaimed_as_card(self):
|
|
found = self.find("iban DE89370400440532013000")
|
|
self.assertNotIn(("370400440532013000", "CREDIT_CARD"), found)
|
|
self.assertEqual({t for _, t in found}, {"BANK_ACCOUNT"})
|
|
|
|
def test_aba_needs_checksum_and_cue(self):
|
|
self.assertIn(("021000021", "BANK_ACCOUNT"), self.find("routing number 021000021"))
|
|
self.assertIn(("021000021", "BANK_ACCOUNT"), self.find("ABA: 021000021"))
|
|
self.assertEqual(self.find("order 021000021 shipped"), set()) # no cue
|
|
self.assertEqual(self.find("routing number 021000022"), set()) # bad checksum
|
|
|
|
def test_bitcoin_base58check(self):
|
|
good = "1BvBMSEYstWetqTFn5Au4m4GFg7xJaNVN2"
|
|
self.assertIn((good, "BANK_ACCOUNT"), self.find("send to %s" % good))
|
|
self.assertIn(("1A1zP1eP5QGefi2DMPTfTL5SLmv7DivfNa", "BANK_ACCOUNT"),
|
|
self.find("genesis 1A1zP1eP5QGefi2DMPTfTL5SLmv7DivfNa"))
|
|
self.assertEqual(self.find("send to 1BvBMSEYstWetqTFn5Au4m4GFg7xJaNVN3"), set())
|
|
|
|
def test_secrets_are_credentials(self):
|
|
cases = {
|
|
"AKIAIOSFODNN7EXAMPLE": "aws key AKIAIOSFODNN7EXAMPLE",
|
|
"ghp_" + "a" * 36: "token ghp_" + "a" * 36 + " end",
|
|
"xoxb-1234567890-abcdefghij": "slack xoxb-1234567890-abcdefghij",
|
|
"AIza" + "b" * 35: "gkey AIza" + "b" * 35,
|
|
"sk_live_" + "c" * 24: "stripe sk_live_" + "c" * 24,
|
|
"glpat-" + "d" * 20: "gl glpat-" + "d" * 20,
|
|
"eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U":
|
|
"bearer eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U",
|
|
}
|
|
for value, text in cases.items():
|
|
self.assertIn((value, "CREDENTIAL"), self.find(text), value)
|
|
|
|
def test_pem_block_is_one_credential(self):
|
|
pem = "-----BEGIN EC PRIVATE KEY-----\nMHQCAQEEIF12 3456 7890 1234\n-----END EC PRIVATE KEY-----"
|
|
found = self.find("key:\n" + pem + "\nok")
|
|
self.assertEqual(found, {(pem, "CREDENTIAL")})
|
|
|
|
def test_secrets_do_not_fire_on_prose(self):
|
|
self.assertEqual(self.find("the skate park (AKA the bowl) opens at 9"), set())
|
|
|
|
|
|
class VaultTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.vault = Vault(os.path.join(self.tmp.name, "v.db"))
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_stable_token_same_scope(self):
|
|
t1 = self.vault.token_for("s1", "EMAIL", "a@b.com")
|
|
t2 = self.vault.token_for("s1", "EMAIL", "a@b.com")
|
|
self.assertEqual(t1, t2)
|
|
self.assertEqual(t1, "[PII_EMAIL_1]")
|
|
|
|
def test_scopes_isolated(self):
|
|
self.vault.token_for("s1", "EMAIL", "a@b.com")
|
|
self.vault.token_for("s1", "EMAIL", "c@d.com")
|
|
t_other = self.vault.token_for("s2", "EMAIL", "c@d.com")
|
|
self.assertEqual(t_other, "[PII_EMAIL_1]") # own sequence
|
|
self.assertIsNone(self.vault.value_for("s2", "[PII_EMAIL_2]"))
|
|
|
|
def test_inventory_has_no_values(self):
|
|
self.vault.token_for("s1", "PERSON", "Ada Lovelace")
|
|
inv = json.dumps(self.vault.scope_tokens("s1"))
|
|
self.assertNotIn("Ada", inv)
|
|
|
|
|
|
class PipelineTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.vault = Vault(os.path.join(self.tmp.name, "v.db"))
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_grounding_drops_hallucinated_span(self):
|
|
text = "call me at 415-555-0173"
|
|
fake = lambda t: [("415-555-0173", "PHONE"), ("Bob Jones", "PERSON")]
|
|
out, res = pipeline.sterilize(text, "s", self.vault, fake)
|
|
self.assertEqual(res["stats"]["llm_ungrounded"], 1)
|
|
self.assertNotIn("Bob", out)
|
|
self.assertNotIn("415-555-0173", out)
|
|
|
|
def test_longest_value_wins(self):
|
|
text = "John Smith emailed. John was brief."
|
|
fake = lambda t: [("John Smith", "PERSON"), ("John", "PERSON")]
|
|
out, _ = pipeline.sterilize(text, "s", self.vault, fake)
|
|
self.assertEqual(out, "[PII_PERSON_1] emailed. [PII_PERSON_2] was brief.")
|
|
|
|
def test_non_pii_bytes_identical(self):
|
|
text = "The nftables ruleset, with odd spacing\tand café unicode, is fine."
|
|
out, res = pipeline.sterilize(text, "s", self.vault, lambda t: [])
|
|
self.assertEqual(out, text)
|
|
self.assertEqual(res["findings"], [])
|
|
|
|
def test_rules_catch_what_llm_missed(self):
|
|
text = "ssn is 123-45-6789 ok"
|
|
out, res = pipeline.sterilize(text, "s", self.vault, lambda t: [])
|
|
self.assertEqual(out, "ssn is [PII_SSN_1] ok")
|
|
self.assertEqual(res["stats"]["rules"], 1)
|
|
|
|
def test_round_trip(self):
|
|
text = "Sam Rivera (sam.r@example.com, 415-555-0173) lives at 4012 Maple Court."
|
|
fake = lambda t: [("Sam Rivera", "PERSON"), ("4012 Maple Court", "ADDRESS")]
|
|
out, _ = pipeline.sterilize(text, "s", self.vault, fake)
|
|
self.assertNotIn("Maple", out)
|
|
self.assertNotIn("sam.r", out)
|
|
back, res = pipeline.rehydrate(out, "s", self.vault)
|
|
self.assertEqual(back, text)
|
|
self.assertEqual(res["unknown_tokens"], [])
|
|
|
|
def test_rehydrate_unknown_token_left_alone(self):
|
|
out, res = pipeline.rehydrate("hello [PII_EMAIL_9]", "s", self.vault)
|
|
self.assertEqual(out, "hello [PII_EMAIL_9]")
|
|
self.assertEqual(res["unknown_tokens"], ["[PII_EMAIL_9]"])
|
|
|
|
def test_stable_across_calls(self):
|
|
fake = lambda t: []
|
|
out1, _ = pipeline.sterilize("mail a@b.com", "s", self.vault, fake)
|
|
out2, _ = pipeline.sterilize("again a@b.com!", "s", self.vault, fake)
|
|
self.assertIn("[PII_EMAIL_1]", out1)
|
|
self.assertIn("[PII_EMAIL_1]", out2)
|
|
|
|
def test_rules_type_wins_on_duplicate_value(self):
|
|
# LLM mistyping a value the rules also caught must not change the
|
|
# minted type — rules are the higher-precision typer.
|
|
text = "reach me at 415-555-0173"
|
|
fake = lambda t: [("415-555-0173", "ID_NUMBER")]
|
|
_, res = pipeline.sterilize(text, "s", self.vault, fake)
|
|
self.assertEqual(res["findings"], [{"token": "[PII_PHONE_1]", "type": "PHONE"}])
|
|
|
|
def assertTokenizedAndRoundTrips(self, text, values, secret, verify=None, label=""):
|
|
vault = Vault(os.path.join(self.tmp.name, "case-%d.db" % self._n))
|
|
self._n += 1
|
|
out, res = pipeline.sterilize(text, "s", vault, lambda t: list(values), verify)
|
|
self.assertNotIn(secret, out, "%s leaked: %r" % (label, out))
|
|
back, _ = pipeline.rehydrate(out, "s", vault)
|
|
self.assertEqual(back, text, "%s round trip: %r" % (label, out))
|
|
return out, res
|
|
|
|
# A value is replaced as a plain substring, inside longer words too:
|
|
# an inflection, a file name or a handle that keeps the value is still
|
|
# the value. Each case must tokenize and rehydrate byte-identically.
|
|
SUBSTRING_CASES = [
|
|
("possessive", "Marisol's file is here", [("Marisol", "PERSON")], "Marisol"),
|
|
("possessive curly", "Marisol\u2019s file", [("Marisol", "PERSON")], "Marisol"),
|
|
("hyphen", "Marisol-Garcia signed", [("Marisol", "PERSON")], "Marisol"),
|
|
("punctuation", "Ask Marisol. Then Marisol, Marisol!", [("Marisol", "PERSON")], "Marisol"),
|
|
("line start and end", "Marisol\nhi\nbye Marisol", [("Marisol", "PERSON")], "Marisol"),
|
|
("email local part", "write marisol.garcia@acme.org or Marisol", [("Marisol", "PERSON")], "Marisol"),
|
|
("url path", "see https://acme.org/team/Marisol/cv", [("Marisol", "PERSON")], "Marisol"),
|
|
("underscore file name", "attached Marisol_Garcia_resume.pdf", [("Marisol", "PERSON")], "Marisol"),
|
|
("underscore username", "login marisol_g from home", [("marisol", "USERNAME")], "marisol"),
|
|
("digit suffix username", "user jsmith42 logged in", [("jsmith", "USERNAME")], "jsmith"),
|
|
("camelCase", "folder MarisolGarcia/", [("Marisol", "PERSON")], "Marisol"),
|
|
("plural surname", "the Garcias came over", [("Garcia", "PERSON")], "Garcia"),
|
|
("inside another word", "Ann flew to Annapolis", [("Ann", "PERSON")], "Ann"),
|
|
("german genitive", "Annas Auto ist rot", [("Anna", "PERSON")], "Anna"),
|
|
("finnish case", "Kerro Marisolille", [("Marisol", "PERSON")], "Marisol"),
|
|
("hungarian case", "Marisolnak adtam", [("Marisol", "PERSON")], "Marisol"),
|
|
("hashtag and digits", "#Marisol2024 party", [("Marisol", "PERSON")], "Marisol"),
|
|
("diacritics", "Jos\u00e9 and Zo\u00eb's", [("Jos\u00e9", "PERSON"), ("Zo\u00eb", "PERSON")], "Jos\u00e9"),
|
|
("decomposed accent", "Jose\u0301 came", [("Jose", "PERSON")], "Jose"),
|
|
("apostrophe name", "O'Brien's O'Brien", [("O'Brien", "PERSON")], "O'Brien"),
|
|
("value with digit", "Room 101B and apt 101", [("101B", "ADDRESS")], "101B"),
|
|
("next to a placeholder", "m@x.com Marisol m@x.comMarisol", [("Marisol", "PERSON")], "Marisol"),
|
|
("chinese", "\u5f20\u4f1f\u4f4f\u5728\u5317\u4eac\uff0c\u5f20\u4f1f\u660e\u4e5f\u5728",
|
|
[("\u5f20\u4f1f", "PERSON")], "\u5f20\u4f1f"),
|
|
("latin next to chinese", "Ann\u4f4f\u5728\u5317\u4eac Ann's", [("Ann", "PERSON")], "Ann"),
|
|
("halfwidth katakana", "\uff80\uff85\uff76\u3055\u3093 \uff80\uff85\uff76\uff80\uff9b\uff73",
|
|
[("\uff80\uff85\uff76", "PERSON")], "\uff80\uff85\uff76"),
|
|
("thai", "\u0e2a\u0e21\u0e0a\u0e32\u0e22\u0e44\u0e1b\u0e1a\u0e49\u0e32\u0e19",
|
|
[("\u0e2a\u0e21\u0e0a\u0e32\u0e22", "PERSON")], "\u0e2a\u0e21\u0e0a\u0e32\u0e22"),
|
|
("korean particle", "\ubbfc\uc218\ub294 \ud559\uad50\uc5d0", [("\ubbfc\uc218", "PERSON")], "\ubbfc\uc218"),
|
|
("devanagari suffix", "\u0930\u093e\u092e\u0915\u094b \u0918\u0930",
|
|
[("\u0930\u093e\u092e", "PERSON")], "\u0930\u093e\u092e"),
|
|
("tai tham", "\u1a20\u1a21\u1a22\u1a23", [("\u1a20\u1a21", "PERSON")], "\u1a20\u1a21"),
|
|
("model value partly overlaps a rule hit", "call 415-555-0173 Main St",
|
|
[("0173 Main St", "ADDRESS")], "Main St"),
|
|
("model value covers a rule hit", "Marisol <m@example.com>",
|
|
[("Marisol <m@example.com>", "OTHER")], "m@example.com"),
|
|
# rule findings: every copy of the value goes, not only the detected one
|
|
("repeated phone", "Call 415-555-0173 or text 415-555-0173.", [], "415-555-0173"),
|
|
("repeated ip, sentence end", "ip 10.1.2.3 and gateway was 10.1.2.3.", [], "10.1.2.3"),
|
|
("ssn copy with suffix", "ssn 123-45-6789 / ref 123-45-6789-A", [], "123-45-6789"),
|
|
("phone copy glued to a word", "tel:415-555-0173 ph415-555-0173", [], "415-555-0173"),
|
|
("phone at sentence end", "Call me at 415-555-0173.", [], "415-555-0173"),
|
|
]
|
|
|
|
def test_substring_cases_tokenize_and_round_trip(self):
|
|
self._n = 0
|
|
for label, text, values, secret in self.SUBSTRING_CASES:
|
|
with self.subTest(label):
|
|
self.assertTokenizedAndRoundTrips(text, values, secret, label=label)
|
|
|
|
def test_substring_match_is_counted_in_findings(self):
|
|
self._n = 0
|
|
_, res = self.assertTokenizedAndRoundTrips(
|
|
"the Garcias came over", [("Garcia", "PERSON")], "Garcia")
|
|
self.assertEqual(res["findings"], [{"token": "[PII_PERSON_1]", "type": "PERSON"}])
|
|
|
|
def test_partial_overlap_becomes_one_token(self):
|
|
# the model value overlaps the rule hit from the right, then from the
|
|
# left; either way one token over the union, rehydrating to the bytes
|
|
cases = [
|
|
("call 415-555-0173 Main St", [("0173 Main St", "ADDRESS")],
|
|
"415-555-0173 Main St", "call [PII_PHONE_1]"),
|
|
("at Main St 415-555-0173 today", [("Main St 415", "ADDRESS")],
|
|
"Main St 415-555-0173", "at [PII_PHONE_1] today"),
|
|
]
|
|
for i, (text, values, merged, want) in enumerate(cases):
|
|
with self.subTest(text):
|
|
vault = Vault(os.path.join(self.tmp.name, "po-%d.db" % i))
|
|
out, _ = pipeline.sterilize(text, "s", vault, lambda t, v=values: v)
|
|
self.assertEqual(out, want)
|
|
self.assertEqual(vault.value_for("s", "[PII_PHONE_1]"), merged)
|
|
back, res = pipeline.rehydrate(out, "s", vault)
|
|
self.assertEqual(back, text)
|
|
self.assertEqual((res["restored"], res["unknown_tokens"]), (1, []))
|
|
|
|
def test_self_overlapping_value_leaves_no_fragment(self):
|
|
for text, value, want in (("x ababa x", "aba", "x [PII_PERSON_1] x"),
|
|
("id 121212 end", "1212", "id [PII_ID_NUMBER_1] end")):
|
|
with self.subTest(text):
|
|
vault = Vault(os.path.join(self.tmp.name, "so-%s.db" % value))
|
|
typ = "ID_NUMBER" if value.isdigit() else "PERSON"
|
|
out, _ = pipeline.sterilize(text, "s", vault, lambda t, v=value, ty=typ: [(v, ty)])
|
|
self.assertEqual(out, want)
|
|
back, _ = pipeline.rehydrate(out, "s", vault)
|
|
self.assertEqual(back, text)
|
|
|
|
def test_never_minted_placeholder_in_the_source_survives_the_round_trip(self):
|
|
# placeholder-shaped text the vault does not hold is escaped, so it
|
|
# comes back as written, card-like digits inside it included
|
|
pipeline.sterilize("Marisol called", "s", self.vault, lambda t: [("Marisol", "PERSON")])
|
|
text = "log: [PII_EMAIL_7], [PII_FOO_99] and [PII_A_4111111111111111], sent by Ana"
|
|
out, res = pipeline.sterilize(text, "s", self.vault, lambda t: [("Ana", "PERSON")])
|
|
for lit in ("[PII_EMAIL_7]", "[PII_FOO_99]", "[PII_A_4111111111111111]"):
|
|
self.assertNotIn(lit, out)
|
|
self.assertEqual(res["stats"]["literal_tokens_escaped"], 3)
|
|
self.assertEqual(res["findings"], [{"token": "[PII_PERSON_2]", "type": "PERSON"}])
|
|
back, info = pipeline.rehydrate(out, "s", self.vault)
|
|
self.assertEqual(back, text)
|
|
self.assertEqual(info["unknown_tokens"], [])
|
|
|
|
def test_live_token_in_the_source_passes_through(self):
|
|
# a placeholder the vault holds for this scope is treated as the token
|
|
# it names: a quoted "[PII_PERSON_1]" rehydrates to the vault value.
|
|
# That is the documented limit; it is what keeps re-sterilizing
|
|
# already-sterilized text stable.
|
|
pipeline.sterilize("Marisol called", "s", self.vault, lambda t: [("Marisol", "PERSON")])
|
|
text = "the log shows [PII_PERSON_1]"
|
|
out, res = pipeline.sterilize(text, "s", self.vault, lambda t: [])
|
|
self.assertEqual(out, text)
|
|
self.assertEqual(res["stats"]["literal_tokens_escaped"], 0)
|
|
self.assertEqual(pipeline.rehydrate(out, "s", self.vault)[0], "the log shows Marisol")
|
|
|
|
def test_escape_held_literals_escapes_a_live_placeholder(self):
|
|
# raw documents: a quoted placeholder is never another text's value
|
|
pipeline.sterilize("mail ana@x.org", "s", self.vault, lambda t: [])
|
|
text = "quoted [PII_EMAIL_1] here"
|
|
kept, _ = pipeline.sterilize(text, "s", self.vault, lambda t: [])
|
|
self.assertEqual(kept, text) # default: the live token
|
|
out, res = pipeline.sterilize(text, "s", self.vault, lambda t: [], escape_held_literals=True)
|
|
self.assertNotIn("[PII_EMAIL_1]", out)
|
|
self.assertEqual(res["stats"]["literal_tokens_escaped"], 1)
|
|
self.assertEqual(pipeline.rehydrate(out, "s", self.vault)[0], text)
|
|
|
|
def test_each_distinct_placeholder_is_looked_up_once(self):
|
|
text = " ".join(["[PII_EMAIL_%d]" % (i % 3) for i in range(300)])
|
|
calls = []
|
|
real = self.vault.value_for
|
|
|
|
def counting(scope, token):
|
|
calls.append(token)
|
|
return real(scope, token)
|
|
|
|
with mock.patch.object(self.vault, "value_for", counting):
|
|
pipeline.sterilize(text, "s", self.vault, lambda t: [])
|
|
self.assertEqual(sorted(calls), ["[PII_EMAIL_0]", "[PII_EMAIL_1]", "[PII_EMAIL_2]"])
|
|
|
|
def test_resterilizing_stored_text_is_stable(self):
|
|
# a stored conversation is sent back to sterilize on every save:
|
|
# three load/append/save rounds, then one rehydrate gives the original
|
|
scope = "thread:abc"
|
|
msgs = [{"role": "user", "content": "Email Marisol at marisol@acme.org"}]
|
|
model = lambda t: [("Marisol", "PERSON")] if "Marisol" in t else []
|
|
stored, _ = pipeline.sterilize(json.dumps(msgs), scope, self.vault, model)
|
|
rows_after_first_save = len(self.vault.scope_tokens(scope))
|
|
for i in range(3):
|
|
loaded = json.loads(stored)
|
|
loaded.append({"role": "assistant", "content": "ok %d" % i})
|
|
msgs.append({"role": "assistant", "content": "ok %d" % i})
|
|
stored, res = pipeline.sterilize(json.dumps(loaded), scope, self.vault, model)
|
|
self.assertEqual(res["stats"]["literal_tokens_escaped"], 0)
|
|
self.assertEqual(json.loads(stored)[0]["content"], "Email [PII_PERSON_1] at [PII_EMAIL_1]")
|
|
self.assertEqual(len(self.vault.scope_tokens(scope)), rows_after_first_save)
|
|
back, _ = pipeline.rehydrate(stored, scope, self.vault)
|
|
self.assertEqual(json.loads(back), msgs)
|
|
|
|
def test_touching_matches_stay_two_tokens(self):
|
|
# adjacent is not overlapping: each value keeps its own token
|
|
text = "MarisolLopez and LopezMarisol"
|
|
out, res = pipeline.sterilize(text, "s", self.vault,
|
|
lambda t: [("Marisol", "PERSON"), ("Lopez", "PERSON")])
|
|
self.assertEqual(out, "[PII_PERSON_1][PII_PERSON_2] and [PII_PERSON_2][PII_PERSON_1]")
|
|
self.assertEqual(len(res["findings"]), 2)
|
|
back, _ = pipeline.rehydrate(out, "s", self.vault)
|
|
self.assertEqual(back, text)
|
|
|
|
def test_verify_value_cannot_break_a_token(self):
|
|
# a verify pass reading sterilized text may propose pieces of a
|
|
# placeholder; none may split one, and the round trip must hold
|
|
self._n = 0
|
|
for piece in ("PII_PERSON_1", "PERSON", "PII", "_1"):
|
|
with self.subTest(piece):
|
|
out, _ = self.assertTokenizedAndRoundTrips(
|
|
"Marisol called", [("Marisol", "PERSON")], "Marisol",
|
|
verify=lambda t, piece=piece: [(piece, "OTHER")], label=piece)
|
|
self.assertEqual(out, "[PII_PERSON_1] called")
|
|
|
|
def test_bare_token_name_is_not_usable(self):
|
|
grounded, dropped = pipeline.ground([("PII_PERSON_1", "OTHER"), ("x PII_EMAIL_2 y", "OTHER")],
|
|
"PII_PERSON_1 x PII_EMAIL_2 y")
|
|
self.assertEqual((grounded, dropped), ([], 2))
|
|
|
|
def test_residuals_are_counted(self):
|
|
# the first rules pass misses the email; the residual loop must
|
|
# tokenize it AND report it, even though the next round is clean
|
|
real = detect.detect
|
|
calls = []
|
|
|
|
def first_pass_blind(text):
|
|
calls.append(text)
|
|
return [] if len(calls) == 1 else real(text)
|
|
|
|
with mock.patch.object(detect, "detect", first_pass_blind):
|
|
out, res = pipeline.sterilize("mail a@b.com and c@d.com", "s", self.vault, lambda t: [])
|
|
self.assertEqual(out, "mail [PII_EMAIL_1] and [PII_EMAIL_2]")
|
|
self.assertEqual(res["stats"]["residual_tokenized"], 2)
|
|
self.assertEqual(res["stats"]["rules"], 0)
|
|
|
|
def test_residuals_summed_across_rounds(self):
|
|
# two residual rounds that each find something: the stat is their
|
|
# sum, not the last round's count
|
|
real = detect.detect
|
|
calls = []
|
|
|
|
def one_at_a_time(text):
|
|
calls.append(text)
|
|
if len(calls) == 1:
|
|
return []
|
|
return real(text)[:1]
|
|
|
|
with mock.patch.object(detect, "detect", one_at_a_time):
|
|
out, res = pipeline.sterilize("mail a@b.com and c@d.com", "s", self.vault, lambda t: [])
|
|
self.assertEqual(out, "mail [PII_EMAIL_1] and [PII_EMAIL_2]")
|
|
self.assertEqual(res["stats"]["residual_tokenized"], 2)
|
|
|
|
def test_no_residuals_reports_zero(self):
|
|
_, res = pipeline.sterilize("ssn 123-45-6789", "s", self.vault, lambda t: [])
|
|
self.assertEqual(res["stats"]["residual_tokenized"], 0)
|
|
|
|
def test_verify_pass_tokenizes_grounded_hits(self):
|
|
text = "the patient is Marisol"
|
|
fake_extract = lambda t: []
|
|
fake_verify = lambda t: [("Marisol", "PERSON")]
|
|
out, res = pipeline.sterilize(text, "s", self.vault, fake_extract, fake_verify)
|
|
self.assertEqual(out, "the patient is [PII_PERSON_1]")
|
|
self.assertEqual(res["stats"]["verify_tokenized"], 1)
|
|
|
|
def test_verify_counts_only_what_it_replaced(self):
|
|
# grounded but replacing nothing: a piece of an existing token, and a
|
|
# value already tokenized by the first pass
|
|
fake_verify = lambda t: [("PERSON", "OTHER"), ("Marisol", "PERSON"), ("Lopez", "PERSON")]
|
|
out, res = pipeline.sterilize("Marisol Lopez called", "s", self.vault,
|
|
lambda t: [("Marisol", "PERSON")], fake_verify)
|
|
self.assertEqual(out, "[PII_PERSON_1] [PII_PERSON_2] called")
|
|
self.assertEqual(res["stats"]["verify_tokenized"], 1)
|
|
|
|
|
|
class PromptContractTests(unittest.TestCase):
|
|
def test_prompts_state_full_type_vocabulary(self):
|
|
# The grammar prefix-coerces free-text labels to enum members
|
|
# ("School" -> SSN, observed live). The prompts must therefore
|
|
# state every legal label verbatim; this guards enum/prompt drift.
|
|
from . import llmclient
|
|
for prompt in (llmclient.EXTRACT_SYSTEM, llmclient.VERIFY_SYSTEM):
|
|
for typ in detect.PII_TYPES:
|
|
self.assertIn(typ, prompt)
|
|
|
|
|
|
class _FakeGliner:
|
|
"""Stands in for gliner.GLiNER: returns canned entities, records labels."""
|
|
entities = []
|
|
calls = []
|
|
|
|
def predict_entities(self, text, labels, threshold=0.5):
|
|
_FakeGliner.calls.append((labels, threshold))
|
|
return list(self.entities)
|
|
|
|
|
|
class RecallBackendTests(unittest.TestCase):
|
|
def setUp(self):
|
|
from . import recall
|
|
self.recall = recall
|
|
_FakeGliner.entities, _FakeGliner.calls = [], []
|
|
|
|
def test_default_is_llm(self):
|
|
b = self.recall.build_backend({"LLM_URL": "http://x:1"})
|
|
self.assertEqual(b.name, "llm")
|
|
|
|
def test_unknown_backend_refused(self):
|
|
with self.assertRaises(self.recall.RecallError):
|
|
self.recall.build_backend({"PII_RECALL": "cloud"})
|
|
|
|
def test_gliner_missing_package_fails_closed(self):
|
|
def loader(_):
|
|
raise ImportError("No module named 'gliner'")
|
|
with self.assertRaises(self.recall.RecallError) as ctx:
|
|
self.recall.build_backend({"PII_RECALL": "gliner"}, loader=loader)
|
|
self.assertIn("pip install gliner", str(ctx.exception))
|
|
|
|
def test_gliner_maps_labels_to_vocabulary_and_dedupes(self):
|
|
_FakeGliner.entities = [
|
|
{"text": "Marisol", "label": "person name", "score": 0.9},
|
|
{"text": "Marisol", "label": "person name", "score": 0.8},
|
|
{"text": "Lincoln Elementary", "label": "employer or school or organization"},
|
|
{"text": "junk", "label": "not a label we asked for"},
|
|
]
|
|
b = self.recall.build_backend(
|
|
{"PII_RECALL": "gliner", "GLINER_THRESHOLD": "0.6"}, loader=lambda _: _FakeGliner())
|
|
out = b.extract("Marisol goes to Lincoln Elementary")
|
|
self.assertEqual(out, [("Marisol", "PERSON"), ("Lincoln Elementary", "ORG")])
|
|
labels, threshold = _FakeGliner.calls[-1]
|
|
self.assertEqual(set(labels), set(self.recall.GLINER_LABELS))
|
|
self.assertEqual(threshold, 0.6)
|
|
for typ in set(self.recall.GLINER_LABELS.values()):
|
|
self.assertIn(typ, detect.PII_TYPES)
|
|
|
|
def test_gliner_verify_uses_lower_threshold(self):
|
|
b = self.recall.build_backend({"PII_RECALL": "gliner"}, loader=lambda _: _FakeGliner())
|
|
b.verify("x")
|
|
self.assertEqual(_FakeGliner.calls[-1][1], self.recall.DEFAULT_GLINER_VERIFY_THRESHOLD)
|
|
|
|
def test_gliner_prediction_failure_is_recall_error(self):
|
|
class Broken:
|
|
def predict_entities(self, *a, **k):
|
|
raise RuntimeError("cuda gone")
|
|
b = self.recall.build_backend({"PII_RECALL": "gliner"}, loader=lambda _: Broken())
|
|
with self.assertRaises(self.recall.RecallError):
|
|
b.extract("x")
|
|
|
|
def test_llm_error_is_a_recall_error(self):
|
|
from . import llmclient
|
|
self.assertTrue(issubclass(llmclient.LlmError, self.recall.RecallError))
|
|
|
|
def test_gliner_backend_through_pipeline(self):
|
|
_FakeGliner.entities = [{"text": "Marisol", "label": "person name"},
|
|
{"text": "Ghost", "label": "person name"}]
|
|
b = self.recall.build_backend({"PII_RECALL": "gliner"}, loader=lambda _: _FakeGliner())
|
|
tmp = tempfile.TemporaryDirectory()
|
|
try:
|
|
vault = Vault(os.path.join(tmp.name, "v.db"))
|
|
out, res = pipeline.sterilize("Marisol called 415-555-0173", "s", vault, b.extract)
|
|
finally:
|
|
tmp.cleanup()
|
|
self.assertEqual(out, "[PII_PERSON_1] called [PII_PHONE_1]")
|
|
self.assertEqual(res["stats"]["llm_ungrounded"], 1) # "Ghost" was not in the text
|
|
|
|
|
|
class _StubLlm(BaseHTTPRequestHandler):
|
|
"""Speaks llama.cpp's chat-completions shape; extraction canned per test."""
|
|
findings = []
|
|
|
|
def do_POST(self):
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
self.rfile.read(length)
|
|
content = json.dumps({"findings": self.findings})
|
|
body = json.dumps(
|
|
{"choices": [{"message": {"content": content}}]}
|
|
).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def log_message(self, *a):
|
|
pass
|
|
|
|
|
|
class ServiceTests(unittest.TestCase):
|
|
KEY = "test-key-123"
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.stub = ThreadingHTTPServer(("127.0.0.1", 0), _StubLlm)
|
|
threading.Thread(target=cls.stub.serve_forever, daemon=True).start()
|
|
cls.tmp = tempfile.TemporaryDirectory()
|
|
cfg = service.Config({
|
|
"PII_BIND": "127.0.0.1:0",
|
|
"LLM_URL": f"http://127.0.0.1:{cls.stub.server_address[1]}",
|
|
"PII_DB": os.path.join(cls.tmp.name, "v.db"),
|
|
"PII_API_KEY": cls.KEY,
|
|
})
|
|
cls.srv = service.build_server(cfg)
|
|
cls.port = cls.srv.server_address[1]
|
|
threading.Thread(target=cls.srv.serve_forever, daemon=True).start()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls.srv.shutdown()
|
|
cls.stub.shutdown()
|
|
cls.tmp.cleanup()
|
|
|
|
def call(self, path, payload=None, key=None, method=None):
|
|
url = f"http://127.0.0.1:{self.port}{path}"
|
|
data = json.dumps(payload).encode() if payload is not None else None
|
|
req = urllib.request.Request(url, data=data, method=method)
|
|
req.add_header("Content-Type", "application/json")
|
|
if key:
|
|
req.add_header("Authorization", f"Bearer {key}")
|
|
with urllib.request.urlopen(req, timeout=10) as r:
|
|
return r.status, json.load(r)
|
|
|
|
def test_health_no_auth(self):
|
|
status, body = self.call("/health")
|
|
self.assertEqual((status, body["ok"]), (200, True))
|
|
|
|
def test_auth_required(self):
|
|
with self.assertRaises(urllib.error.HTTPError) as ctx:
|
|
self.call("/sterilize", {"text": "x"}, key="wrong")
|
|
self.assertEqual(ctx.exception.code, 401)
|
|
|
|
def test_sterilize_rehydrate_over_http(self):
|
|
_StubLlm.findings = [{"text": "Marisol", "type": "PERSON"}]
|
|
original = "Marisol's ssn is 123-45-6789, mail her at m@example.com"
|
|
_, st = self.call("/sterilize", {"text": original, "scope": "conv1"}, key=self.KEY)
|
|
for leak in ("Marisol", "123-45-6789", "m@example.com"):
|
|
self.assertNotIn(leak, st["text"])
|
|
self.assertEqual(st["stats"]["rules"], 2)
|
|
_, back = self.call("/rehydrate", {"text": st["text"], "scope": "conv1"}, key=self.KEY)
|
|
self.assertEqual(back["text"], original)
|
|
# wrong scope must not resolve
|
|
_, other = self.call("/rehydrate", {"text": st["text"], "scope": "conv2"}, key=self.KEY)
|
|
self.assertEqual(other["restored"], 0)
|
|
|
|
def test_flag_rules_only_no_vault_rows(self):
|
|
_StubLlm.findings = [{"text": "should never be called", "type": "PERSON"}]
|
|
_, body = self.call("/flag", {"text": "ssn 123-45-6789 and a@b.com"}, key=self.KEY)
|
|
self.assertTrue(body["flagged"])
|
|
self.assertEqual(body["counts"], {"SSN": 1, "EMAIL": 1})
|
|
_, clean = self.call("/flag", {"text": "nothing sensitive here"}, key=self.KEY)
|
|
self.assertEqual((clean["flagged"], clean["counts"]), (False, {}))
|
|
# no scope was touched, nothing minted
|
|
_, inv = self.call("/scope/default", key=self.KEY)
|
|
self.assertEqual(inv["tokens"], [])
|
|
|
|
def test_escape_held_literals_over_http(self):
|
|
_StubLlm.findings = []
|
|
self.call("/sterilize", {"text": "mail a@b.com", "scope": "esc"}, key=self.KEY)
|
|
quote = "quoted [PII_EMAIL_1]"
|
|
_, kept = self.call("/sterilize", {"text": quote, "scope": "esc"}, key=self.KEY)
|
|
self.assertEqual(kept["text"], quote) # default: a re-save keeps the token
|
|
_, esc = self.call("/sterilize", {"text": quote, "scope": "esc",
|
|
"escape_held_literals": True}, key=self.KEY)
|
|
self.assertNotIn("[PII_EMAIL_1]", esc["text"])
|
|
_, back = self.call("/rehydrate", {"text": esc["text"], "scope": "esc"}, key=self.KEY)
|
|
self.assertEqual(back["text"], quote)
|
|
# the rule applied is echoed, so a caller can tell it was honoured
|
|
self.assertIs(kept["stats"]["escape_held_literals"], False)
|
|
self.assertIs(esc["stats"]["escape_held_literals"], True)
|
|
|
|
def test_scope_inventory_no_values(self):
|
|
_StubLlm.findings = []
|
|
self.call("/sterilize", {"text": "a@b.com", "scope": "inv"}, key=self.KEY)
|
|
_, body = self.call("/scope/inv", key=self.KEY)
|
|
self.assertEqual(body["tokens"][0]["token"], "[PII_EMAIL_1]")
|
|
self.assertNotIn("a@b.com", json.dumps(body))
|
|
|
|
def test_fail_closed_without_key(self):
|
|
with self.assertRaises(SystemExit):
|
|
service.Config({"PII_DB": "x"})
|
|
|
|
def test_fail_closed_on_unstartable_recall_backend(self):
|
|
with self.assertRaises(SystemExit) as ctx:
|
|
service.Config({"PII_DB": "x", "PII_API_KEY": "k", "PII_RECALL": "nope"})
|
|
self.assertIn("PII_RECALL", str(ctx.exception))
|
|
|
|
def test_sterilize_reports_recall_backend(self):
|
|
_StubLlm.findings = []
|
|
_, st = self.call("/sterilize", {"text": "a@b.com", "scope": "rb"}, key=self.KEY)
|
|
self.assertEqual(st["stats"]["recall"], "llm")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|