Spaces:
Running
Running
File size: 5,348 Bytes
9158b0a 2633ee9 9158b0a 2633ee9 9158b0a 2633ee9 9158b0a 553ced5 9158b0a 2633ee9 9158b0a 553ced5 9158b0a 2633ee9 9158b0a 553ced5 9158b0a 553ced5 9158b0a |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
import hashlib
import re
import pytest
from hypothesis import given, settings
from hypothesis import strategies as st
from guardrails_genie.guardrails.secrets_detection import (
DEFAULT_SECRETS_PATTERNS,
REDACTION,
SecretsDetectionGuardrail,
SecretsDetectionResponse,
SecretsDetectionSimpleResponse,
redact,
)
@pytest.fixture
def mock_secrets_guard(monkeypatch):
def _mock_guard(*args, **kwargs):
prompt = kwargs.get("prompt")
return_detected_types = kwargs.get("return_detected_types")
if "safe text" in prompt:
if return_detected_types:
return SecretsDetectionResponse(
contains_secrets=False,
explanation="No secrets detected in the text.",
detected_secrets={},
redacted_text=prompt,
)
else:
return SecretsDetectionSimpleResponse(
contains_secrets=False,
explanation="No secrets detected in the text.",
redacted_text=prompt,
)
else:
if return_detected_types:
return SecretsDetectionResponse(
contains_secrets=True,
explanation="The output contains secrets.",
detected_secrets={"secrets": ["API_KEY"]},
redacted_text="My secret key is [REDACTED:]************[:REDACTED]",
)
else:
return SecretsDetectionSimpleResponse(
contains_secrets=True,
explanation="The output contains secrets.",
redacted_text="My secret key is [REDACTED:]************[:REDACTED]",
)
monkeypatch.setattr(
"guardrails_genie.guardrails.secrets_detection.SecretsDetectionGuardrail.guard",
_mock_guard,
)
def test_redact_partial():
text = "My secret key is ABCDEFGHIJKL"
matches = ["ABCDEFGHIJKL"]
redacted_text = redact(text, matches, REDACTION.REDACT_PARTIAL)
assert redacted_text == "My secret key is [REDACTED:]AB..KL[:REDACTED]"
def test_redact_all():
text = "My secret key is ABCDEFGHIJKL"
matches = ["ABCDEFGHIJKL"]
redacted_text = redact(text, matches, REDACTION.REDACT_ALL)
assert redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"
def test_redact_hash():
text = "My secret key is ABCDEFGHIJKL"
matches = ["ABCDEFGHIJKL"]
hashed_value = hashlib.md5("ABCDEFGHIJKL".encode()).hexdigest()
redacted_text = redact(text, matches, REDACTION.REDACT_HASH)
assert redacted_text == f"My secret key is [REDACTED:]{hashed_value}[:REDACTED]"
def test_redact_no_match():
text = "My secret key is ABCDEFGHIJKL"
matches = ["XYZ"]
redacted_text = redact(text, matches, REDACTION.REDACT_ALL)
assert redacted_text == text
def test_secrets_detection_guardrail_detect_types(mock_secrets_guard):
from guardrails_genie.guardrails.secrets_detection import (
REDACTION,
SecretsDetectionGuardrail,
)
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
prompt = "My secret key is ABCDEFGHIJKL"
result = guardrail.guard(prompt=prompt, return_detected_secrets=True)
assert result.contains_secrets is True
assert result.explanation == "The output contains secrets."
assert result.detected_secrets == {"secrets": ["API_KEY"]}
assert result.redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"
def test_secrets_detection_guardrail_simple_response(mock_secrets_guard):
from guardrails_genie.guardrails.secrets_detection import (
REDACTION,
SecretsDetectionGuardrail,
)
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
prompt = "My secret key is ABCDEFGHIJKL"
result = guardrail.guard(prompt=prompt, return_detected_secrets=False)
assert result.contains_secrets is True
assert result.explanation == "The output contains secrets."
assert result.redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"
def test_secrets_detection_guardrail_no_secrets(mock_secrets_guard):
from guardrails_genie.guardrails.secrets_detection import (
REDACTION,
SecretsDetectionGuardrail,
)
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
prompt = "This is a safe text with no secrets."
result = guardrail.guard(prompt=prompt, return_detected_secrets=True)
assert result.contains_secrets is False
assert result.explanation == "No secrets detected in the text."
assert result.detected_secrets == {}
assert result.redacted_text == prompt
# Create a strategy to generate strings that match the patterns
def pattern_strategy(pattern):
return st.from_regex(re.compile(pattern), fullmatch=True)
@settings(deadline=1000) # Set the deadline to 1000 milliseconds (1 second)
@given(pattern_strategy(DEFAULT_SECRETS_PATTERNS["JwtToken"][0]))
def test_specific_pattern_guardrail(text):
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
result = guardrail.guard(prompt=text, return_detected_secrets=True)
assert result.contains_secrets is True
assert "JwtToken" in result.detected_secrets
|