Spaces:
Configuration error
Configuration error
File size: 1,414 Bytes
447ebeb |
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 |
import json
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
from litellm.types.guardrails import SupportedGuardrailIntegrations
def test_initialize_presidio_guardrail():
"""
Test that initialize_guardrail correctly uses registered initializers
for presidio guardrail
"""
# Setup test data for a non-custom guardrail (using Presidio as an example)
test_guardrail = {
"guardrail_name": "test_presidio_guardrail",
"litellm_params": {
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
"mode": "pre_call",
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
},
}
# Call the initialize_guardrail method
guardrail_handler = InMemoryGuardrailHandler()
result = guardrail_handler.initialize_guardrail(
guardrail=test_guardrail,
)
assert result["guardrail_name"] == "test_presidio_guardrail"
assert (
result["litellm_params"].guardrail
== SupportedGuardrailIntegrations.PRESIDIO.value
)
assert result["litellm_params"].mode == "pre_call"
|