Skip to content

test(processor): add unit tests for response.go's ProcessResponse - #610

Open
kalra-mohit wants to merge 1 commit into
dataiku:mainfrom
kalra-mohit:tests/response-processor-429
Open

test(processor): add unit tests for response.go's ProcessResponse#610
kalra-mohit wants to merge 1 commit into
dataiku:mainfrom
kalra-mohit:tests/response-processor-429

Conversation

@kalra-mohit

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

Copy link
Copy Markdown

Closes #429.

response.go already had decent coverage on RestorePII/BuildRestorer, but ProcessResponse — the function that actually gets called on every response — had none. This adds it.

ProcessResponse's job, stripped down: gate on content type, try to parse the body as JSON (fall back to the original bytes if it doesn't parse), hand off to the provider's RestoreMaskedResponse for the actual PII restoration, then stamp original_response and proxy_metadata onto the result.

What's covered:

  • Non-JSON content type passes through byte-for-byte, untouched.
  • Malformed JSON falls back to the original body instead of erroring or emitting something truncated.
  • The main path, end-to-end: a masked placeholder in choices[0].message.content gets restored to the real value, original_response holds the untouched input bytes, and proxy_metadata gets stamped. original_response is embedded via json.RawMessage, so it comes back as nested JSON rather than a string when you unmarshal the result — took me a second to realize the test needed to account for that instead of comparing it as a plain string.
  • The "mapping miss" case from the issue: if a masked-looking token in the response has no matching entry in maskedToOriginal, it's left exactly as-is — not corrupted, not blanked.
  • What happens when the provider can't make sense of the response shape (RestoreMaskedResponse errors on something without a choices field): ProcessResponse just logs it and keeps going, and proxy_metadata still gets added. That distinction matters — a caller shouldn't be able to tell "nothing needed restoring" apart from "the proxy gave up silently."
  • The proxy-notice wiring: confirms the interception-notice string ProcessResponse builds, and the getAddProxyNotice callback it passes down, actually reach OpenAIProvider and show up in the response content.

Why a real OpenAIProvider instead of a mock: providers.Provider's CreateMaskedRequest and RestoreMaskedResponse are typed against unexported function types declared inside package providers (maskPIIInTextType, restorePIIType, etc.), so nothing defined outside that package can satisfy the interface — a hand-rolled fake from processor_test simply won't compile. I noticed proxy/handler_test.go already deals with this the same way (constructs a real OpenAIProvider, which is cheap since it makes no network calls until you actually send a request), so I followed that existing pattern instead of inventing a different one.

One thing I left deliberately undone: the json.Marshal failure fallback (the branch where re-marshaling the modified response fails and ProcessResponse returns the original body) isn't separately tested. Triggering it needs a provider to write something unmarshalable — a channel, a func — into the decoded response map, and every real provider only ever writes strings. Given the same unexported-interface constraint above, there's no way to force that branch through the public API without a fake, which isn't buildable here. I left a comment explaining that rather than skip it silently.

This repo requires first-time contributors to add themselves to CONTRIBUTORS.md; that entry is in #609 (one of three PRs from the same pass), not duplicated here.

flowchart LR
    A[Incoming response body\n+ contentType + maskedToOriginal] --> B{JSON content type?}
    B -->|no| C[return body unchanged]:::covered
    B -->|yes| D{json.Unmarshal succeeds?}
    D -->|no, malformed| E[return original body]:::covered
    D -->|yes| F["provider.RestoreMaskedResponse\n(e.g. OpenAIProvider)"]
    F -->|error - unrecognized shape| G[log only, keep going]:::covered
    F -->|ok| H["response with PII\nrestored in place\n(unmapped tokens left as-is)"]:::covered
    G --> I[stamp original_response\n+ proxy_metadata]:::covered
    H --> I
    I --> J{json.Marshal succeeds?}
    J -->|yes| K[return modified JSON]:::covered
    J -->|no - not reachable\nvia public API today| L[return original body]:::uncovered

    classDef covered fill:#1f4d2c,stroke:#27ae60,color:#fff
    classDef uncovered stroke-dasharray: 5 5
Loading

Test plan:

  • go test ./src/backend/processor/... passes
  • go vet and gofmt -l clean
  • For the two most consequential assertions — the original_response field and the provider-error-keeps-going behavior — I temporarily removed the field and made the error path return early instead of continuing, confirmed both new tests failed as expected, then reverted

The file already had RestorePII/BuildRestorer coverage, but its main
exported entry point, ProcessResponse, had none. Adds tests for: the
content-type gate (non-JSON bodies pass through untouched), malformed
JSON falling back to the original body, a real end-to-end restore through
a genuine OpenAIProvider (masked value restored, original_response
preserved as raw untouched bytes, proxy_metadata stamped), an unmapped
masked token being left alone rather than corrupted (the "mapping miss"
case), the provider-returns-an-error path still stamping metadata instead
of aborting, and the interception-notice wiring.

providers.Provider can't be faked from this package — CreateMaskedRequest
and RestoreMaskedResponse are typed against unexported function types
declared in package providers, so only types in that package can satisfy
the interface. Tests use a real OpenAIProvider instead (it makes no
network calls until a request is actually sent), matching the existing
convention in proxy/handler_test.go rather than inventing a workaround.

Closes dataiku#429

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@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 added a commit to kalra-mohit/kiji-proxy that referenced this pull request Jul 22, 2026
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

Copy link
Copy Markdown
Author

recheck

@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/processor/response.go

1 participant