50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
|
import unittest
|
||
|
|
from guardrails.telemetry.common import recursive_key_operation, redact
|
||
|
|
import ast
|
||
|
|
|
||
|
|
|
||
|
|
# Test suite for recursive_key_operation function
|
||
|
|
class TestRecursiveKeyOperation(unittest.TestCase):
|
||
|
|
def test_list(self):
|
||
|
|
data = '{"init_args": [], "init_kwargs": {"model": "gpt-4o-mini", \
|
||
|
|
"api_base": "https://api.openai.com/v1", "api_key": "sk-1234"}}'
|
||
|
|
result = recursive_key_operation(data, redact)
|
||
|
|
assert ast.literal_eval(result)["init_kwargs"]["api_key"] == "***1234"
|
||
|
|
|
||
|
|
def test_dict_kwargs(self):
|
||
|
|
data = {
|
||
|
|
"index": "0",
|
||
|
|
"api": '{"init_args": [], "init_kwargs": {"model": "gpt-4o-mini",\
|
||
|
|
"api_base": "https://api.openai.com/v1", "api_key": "sk-1234"}}',
|
||
|
|
"messages": None,
|
||
|
|
"prompt_params": "{}",
|
||
|
|
"output_schema": '{"type": "string"}',
|
||
|
|
"output": None,
|
||
|
|
}
|
||
|
|
result = recursive_key_operation(data, redact)
|
||
|
|
assert ast.literal_eval(result["api"])["init_kwargs"]["api_key"] == "***1234"
|
||
|
|
|
||
|
|
def test_nomatch(self):
|
||
|
|
data = {"somekey": "soemvalue"}
|
||
|
|
result = recursive_key_operation(data, redact)
|
||
|
|
self.assertEqual(result, data)
|
||
|
|
|
||
|
|
def test_empty_dict(self):
|
||
|
|
data = {}
|
||
|
|
result = recursive_key_operation(data, redact)
|
||
|
|
self.assertEqual(result, data)
|
||
|
|
|
||
|
|
def test_empty_list(self):
|
||
|
|
data = []
|
||
|
|
result = recursive_key_operation(data, redact)
|
||
|
|
self.assertEqual(result, data)
|
||
|
|
|
||
|
|
def test_non_string_value(self):
|
||
|
|
data = {"key": 123, "another_key": "value"}
|
||
|
|
result = recursive_key_operation(data, redact)
|
||
|
|
self.assertEqual(result, data)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|