From 7d692eaf9932f00c4ed5f49f675b8f1ced9a0c7e Mon Sep 17 00:00:00 2001 From: Travis Matthews Date: Tue, 13 Oct 2026 16:49:05 -0400 Subject: [PATCH] Make the pii occurrence pagination test independent of numeric labels The test asserted that 12 random "The number is N" values all came back as NUMERIC_VALUE. A newer English NER model labels some bare 3-digit numbers CREDIT_CARD (score > the 0.5 regex NUMERIC_VALUE, so it wins the overlap), which made the count vary with the random numbers even though pagination worked. Compare paginated results against a single unpaginated fetch, and check every number was found under some label, instead. Co-Authored-By: Claude Opus 5.5 --- .../dataset_tests/test_pii_occurrences.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/tests/dataset_tests/test_pii_occurrences.py b/tests/tests/dataset_tests/test_pii_occurrences.py index fb66773..85e6b20 100644 --- a/tests/tests/dataset_tests/test_pii_occurrences.py +++ b/tests/tests/dataset_tests/test_pii_occurrences.py @@ -57,7 +57,7 @@ def test_get_occurences(textual: TextualNer, created_datasets): def test_occurences_pagination(textual: TextualNer, created_datasets): #we fetch 5 records at a time per entity, so lets create a file with 12 numeric values - random_numbers = [random.randint(0, 1000) for _ in range(12)] + random_numbers = random.sample(range(1001), 30) lines = ['The number is ' + str(r) for r in random_numbers] document = '\n'.join(lines) @@ -81,12 +81,24 @@ def test_occurences_pagination(textual: TextualNer, created_datasets): assert len(processed_files)==1, "Failed to process files" file: DatasetFile = list(filter(lambda x: x.name=='file.txt', processed_files))[0] + # Which label a bare number gets (NUMERIC_VALUE, CREDIT_CARD, ...) depends on the model, so + # compare paginated results against a single unpaginated fetch rather than a fixed label count. file._pii_occurence_file_limit=5 + paged = file.get_entities() + paged_mixed = file.get_entities([PiiType.NUMERIC_VALUE, 'NAME_FAMILY']) #we add a string value to make sure we support both str and piitype instances - entities = file.get_entities([PiiType.NUMERIC_VALUE, 'NAME_FAMILY']) #we add a string value to make sure we support both str and piitype instances - occurrences = entities[PiiType.NUMERIC_VALUE] + file._pii_occurence_file_limit=1000 + unpaged = file.get_entities() + unpaged_mixed = file.get_entities([PiiType.NUMERIC_VALUE, 'NAME_FAMILY']) - assert len(occurrences)==12, 'failed pagination' + assert occurrence_keys(paged) == occurrence_keys(unpaged), 'failed pagination' + assert occurrence_keys(paged_mixed) == occurrence_keys(unpaged_mixed), 'failed pagination' + + found = {o['entity'] for occurrences in paged.values() for o in occurrences} + assert {str(r) for r in random_numbers} <= found, 'failed pagination' + +def occurrence_keys(entities): + return sorted((pii_type, o['entity'], o['head'], o['tail']) for pii_type, occurrences in entities.items() for o in occurrences) def create_file_stream(txt: str) -> io.BytesIO: return io.BytesIO(txt.encode('utf-8'))