Skip to content

test(pii): add unit tests for masking_service.go#609

Open
kalra-mohit wants to merge 2 commits into
dataiku:mainfrom
kalra-mohit:tests/masking-service-430
Open

test(pii): add unit tests for masking_service.go#609
kalra-mohit wants to merge 2 commits into
dataiku:mainfrom
kalra-mohit:tests/masking-service-430

Conversation

@kalra-mohit

@kalra-mohit kalra-mohit commented Jul 22, 2026

Copy link
Copy Markdown

Closes #430.

masking_service.go was one of the few files in src/backend/pii without its own test file — database.go, the generators, and the detectors all had coverage, but the thing that actually wires them together (MaskText) didn't. Added src/backend/pii/masking_service_test.go.

The tests use a stub detector rather than the real regex/NER models, since detection itself is already covered under pii/detectors — what's under test here is the orchestration around it: position math, dummy generation/reuse, and the sweep step.

A few of these are worth calling out:

  • Duplicate occurrences. A detector typically reports a repeated PII string once even if it appears three times in the text. There's a sweep step in MaskText that catches the rest, and I wanted a test that fails loudly if that sweep ever gets dropped or short-circuited — without it, the second and third occurrence would go upstream unmasked.
  • Dummy reuse across requests. MaskingService is backed by a persistent PIIMapping (real SQLite, via the existing newTestDB helper), so the same original PII should map to the same dummy across two separate calls. I tested that directly rather than trusting it.
  • Byte vs. rune offsets. Entity positions from the detector are byte offsets, not rune offsets, so I added a Unicode case ("Zoë Müller" next to an emoji) to make sure masking a name near multi-byte characters doesn't corrupt its neighbors.
  • Fail-open on detector error. If GetDetector() or Detect() fails, MaskText currently returns the original text with an empty mapping — no PII gets masked, and there's no signal in the result that detection didn't actually run. That's arguably a problem for a privacy proxy (it fails open, not closed, unlike the disabled-entities path), but it's the existing documented behavior, so I pinned it down with a test rather than silently changing behavior in a "just add tests" PR.

While writing the RestorePII tests, I found something worth flagging separately. MaskingService.RestorePII still does sequential strings.ReplaceAll calls over the mapping — which is exactly the chained-substitution bug that processor.ResponseProcessor.RestorePII used to have, before it was rewritten around BuildRestorer (single-pass strings.Replacer) in be948f86. That fix's own comment explains the failure mode: if one dummy happens to equal another entry's original text, sequential replacement corrupts the output depending on map iteration order.

MaskingService.RestorePII never got the same fix. As far as I can tell from grepping the codebase, nothing in the current request path actually calls it — only ResponseProcessor.RestorePII is wired into production traffic — but it's still an exported method, so anything outside this repo (or future code inside it) could hit the bug. Rather than quietly patch it under the cover of "add tests" or just leave it untested, I added TestMaskingServiceRestorePII_ChainedSubstitutionBug, which pins down the current (broken) output using a swap-cycle case (A's original is B's dummy and vice versa) — chosen specifically because it corrupts under either possible Go map-iteration order, so the test isn't flaky. I ran it 200 times locally to be sure. Left a comment on the method recommending either reusing BuildRestorer here too, or deleting the method if it's genuinely dead — happy to put up that follow-up if wanted.

(One correction while writing this up: I'd originally cited the wrong commit for the BuildRestorer fix — aad146d, which is actually an unrelated PII-generator change from #594. The real one is be948f86, "fix: restore PII in a single pass to prevent chained substitution.")

Also added myself to CONTRIBUTORS.md per the first-time-contributor note in CONTRIBUTING.md — this is one of three PRs from the same pass over the PII/processor/server packages (#609, #610, #611); I only added the entry here to avoid a merge conflict across all three.

flowchart TD
    subgraph Mask["Masking path — MaskText (this PR's main focus)"]
        A[Detector.Detect] --> B{Entities found?}
        B -->|no| C[return original text unchanged]
        B -->|yes| D[filter disabled entity types]
        D --> E[generate/reuse dummy per original\nvia persistent PIIMapping]
        E --> F[sweep: replace ALL occurrences\nof each original, not just the\ndetector-reported one]
        F --> G[MaskedText + MaskedToOriginal]
    end

    subgraph Restore["Restore path — two implementations"]
        H["processor.ResponseProcessor.RestorePII\n(live, in the request path)"] --> I["BuildRestorer: single-pass\nstrings.Replacer"]
        I --> J[correct output,\nfixed in be948f86]

        K["MaskingService.RestorePII\n(exported, but NOT called\nanywhere in the request path)"] --> L["sequential strings.ReplaceAll\nper mapping entry"]
        L --> M["chained substitution:\ncan corrupt output if one\ndummy collides with\nanother original"]
    end

    G -.mapping used by both restorers.-> H
    G -.mapping used by both restorers.-> K

    style M fill:#5a1f1f,stroke:#c0392b,color:#fff
    style J fill:#1f4d2c,stroke:#27ae60,color:#fff
    style K stroke-dasharray: 5 5
Loading

Test plan:

  • go test ./src/backend/pii/... passes (see CONTRIBUTING.md for the make setup-tokenizers step needed before this builds locally)
  • go vet and gofmt -l clean
  • For the sweep test and the disabled-entity test, I temporarily broke the corresponding logic (removed the sweep loop, removed the disabled-entity filter) and confirmed the new test failed, then reverted — wanted to be sure these weren't just asserting something that'd pass regardless
  • Ran TestMaskingServiceRestorePII_ChainedSubstitutionBug 200x locally since it depends on map iteration order; no flakes

@github-actions

github-actions Bot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

All contributors have signed the CLA ✍️ ✅
Posted by the CLA Assistant Lite bot.

kalra-mohit and others added 2 commits July 21, 2026 18:34
Adds src/backend/pii/masking_service_test.go, covering MaskText's core
orchestration: multiple PII types in one string, the duplicate-occurrence
sweep (a single detected entity masking every repeat of that string),
same-dummy reuse across requests via the persistent PIIMapping, PII at
string boundaries, the position-fallback path for invalid entity offsets,
unicode/multi-byte input, disabled-entity passthrough, and the two
detector-error paths.

Also documents a real, currently-untested defect found while writing these
tests: MaskingService.RestorePII still does sequential strings.ReplaceAll
over the mapping and is vulnerable to the same chained-substitution bug
that was fixed in processor.ResponseProcessor.RestorePII via BuildRestorer
(see commit be948f8). It looks unused in the current request path (only
ResponseProcessor.RestorePII is wired in), but it's exported, so a
characterization test pins down current behavior and flags it for a
follow-up fix rather than silently leaving it uncovered.

Closes dataiku#430

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
First-time contribution, per CONTRIBUTING.md's requirement to add
yourself to CONTRIBUTORS.md in your first PR. Adding here since this is
the first of three PRs from this contribution; the other two (dataiku#610,
dataiku#611) reference this instead of duplicating the entry.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@kalra-mohit

Copy link
Copy Markdown
Author

I have read the CLA Document and I hereby sign the CLA


@kalra-mohit

Copy link
Copy Markdown
Author

I have read the CLA Document and I hereby sign the CLA

@kalra-mohit
kalra-mohit marked this pull request as ready for review July 22, 2026 02:30
@kalra-mohit

Copy link
Copy Markdown
Author

@hanneshapke whenever you get a chance to take a look, happy to address any feedback!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add unit tests for src/backend/pii/masking_service.go

1 participant