MultiLLM-Proxy/tests/test_redaction.py at main · ALikesToCode/MultiLLM-Proxy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import unittest

from services.redaction import (
    REDACTED,
    redact_headers,
    redact_payload,
    redact_query_params,
    redact_text,
)


class RedactionTest(unittest.TestCase):
    def test_redact_headers_hides_auth_material(self):
        redacted = redact_headers(
            {
                "Authorization": "Bearer sk-secret123456",
                "x-goog-api-key": "AIza-provider-secret",
                "X-MultiLLM-Api-Key": "proxy-secret",
                "X-PAYMENT": "payment-proof",
                "X-Encryption-Key": "encryption-secret",
                "Content-Type": "application/json",
            }
        )

        self.assertEqual(redacted["Authorization"], REDACTED)
        self.assertEqual(redacted["x-goog-api-key"], REDACTED)
        self.assertEqual(redacted["X-MultiLLM-Api-Key"], REDACTED)
        self.assertEqual(redacted["X-PAYMENT"], REDACTED)
        self.assertEqual(redacted["X-Encryption-Key"], REDACTED)
        self.assertEqual(redacted["Content-Type"], "application/json")

    def test_redact_payload_hides_nested_prompts_and_secrets(self):
        payload = {
            "model": "openai/gpt-4o-mini",
            "messages": [{"role": "user", "content": "private prompt"}],
            "metadata": {
                "api_key": "sk-secret123456",
                "client_secret": "navy-secret",
                "code_verifier": "pkce-verifier",
                "refresh_token": "navy-ort-secret",
                "safe": "visible",
            },
        }

        redacted = redact_payload(payload)

        self.assertEqual(redacted["messages"], REDACTED)
        self.assertEqual(redacted["metadata"]["api_key"], REDACTED)
        self.assertEqual(redacted["metadata"]["client_secret"], REDACTED)
        self.assertEqual(redacted["metadata"]["code_verifier"], REDACTED)
        self.assertEqual(redacted["metadata"]["refresh_token"], REDACTED)
        self.assertEqual(redacted["metadata"]["safe"], "visible")
        self.assertNotIn("private prompt", str(redacted))
        self.assertNotIn("sk-secret", str(redacted))

    def test_redact_query_params_hides_keys(self):
        redacted = redact_query_params(
            {
                "key": "AIza-query-secret",
                "model": "gemini-2.0-flash",
            }
        )

        self.assertEqual(redacted["key"], REDACTED)
        self.assertEqual(redacted["model"], "gemini-2.0-flash")

    def test_redact_text_masks_common_secret_shapes(self):
        redacted = redact_text(
            'Authorization: Bearer sk-secret123456; '
            'Authorization: L402 macaroon:preimage; '
            '{"api_key":"AIza-provider-secret","refresh_token":"navy-ort-secret"}'
        )

        self.assertNotIn("sk-secret", redacted)
        self.assertNotIn("AIza-provider-secret", redacted)
        self.assertNotIn("macaroon:preimage", redacted)
        self.assertNotIn("navy-ort-secret", redacted)
        self.assertIn(REDACTED, redacted)


if __name__ == "__main__":
    unittest.main()