scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

add unit test

Browse source
This commit is contained in:
Moritz Müller 2026-03-27 17:38:50 +01:00
parent e0ef9bc411
commit 2df1be6055
1 changed files with 54 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
tests/test_litellm/litellm_core_utils/test_redact_messages.py
View file

@ -8,7 +8,10 @@ but litellm_params["litellm_metadata"] is None.
import pytest
import litellm
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.litellm_core_utils.redact_messages import (
perform_redaction,
should_redact_message_logging,
)
@pytest.fixture(autouse=True)
@ -143,3 +146,53 @@ class TestShouldRedactMessageLogging:
litellm_metadata=None,
)
assert should_redact_message_logging(details) is True
class TestPerformRedaction:
"""Unit tests for perform_redaction()."""
def test_redact_async_complete_streaming_response(self):
"""Test that async_complete_streaming_response is properly redacted."""
response_obj = litellm.ModelResponse(
choices=[
litellm.Choices(
message=litellm.Message(content="secret content", role="assistant")
)
]
)
model_call_details = {
"messages": [{"role": "user", "content": "hi"}],
"prompt": "hi",
"input": "hi",
"stream": True,
"async_complete_streaming_response": response_obj,
}
result = perform_redaction(model_call_details, result=None)
redacted_response = model_call_details["async_complete_streaming_response"]
assert redacted_response.choices[0].message.content == "redacted-by-litellm"
def test_redact_complete_streaming_response(self):
"""Test that complete_streaming_response is properly redacted."""
response_obj = litellm.ModelResponse(
choices=[
litellm.Choices(
message=litellm.Message(content="secret content", role="assistant")
)
]
)
model_call_details = {
"messages": [{"role": "user", "content": "hi"}],
"prompt": "hi",
"input": "hi",
"stream": True,
"complete_streaming_response": response_obj,
}
result = perform_redaction(model_call_details, result=None)
redacted_response = model_call_details["complete_streaming_response"]
assert redacted_response.choices[0].message.content == "redacted-by-litellm"