Skip to content

Commit b781136

Browse files
committed
Add test to ensure get_standard_logging_object_payload is called exactly once on cache hits
1 parent 5617f35 commit b781136

File tree

1 file changed

+34
-0
lines changed

1 file changed

+34
-0
lines changed

tests/test_litellm/test_logging.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
)
3131
from litellm.integrations.custom_logger import CustomLogger
3232
from litellm.types.utils import StandardLoggingPayload
33+
from litellm.litellm_core_utils.litellm_logging import Logging
3334

3435

3536
class CacheHitCustomLogger(CustomLogger):
@@ -152,3 +153,36 @@ async def test_cache_hit_includes_custom_llm_provider():
152153
# Clean up
153154
litellm.callbacks = original_callbacks
154155
litellm.cache = None
156+
157+
158+
@pytest.mark.asyncio
159+
async def test_unified_handler_calls_get_standard_logging_object_payload_once(mocker, monkeypatch):
160+
"""
161+
Tests that for a cache hit,
162+
the test_unified_handler_calls_get_standard_logging_object_payload_once is called exactly once.
163+
"""
164+
monkeypatch.setattr(litellm, "cache", litellm.Cache())
165+
test_message = [{"role": "user", "content": "helper function call test"}]
166+
167+
mock_helper_fn = mocker.spy(Logging, "_success_handler_helper_fn")
168+
mock_get_standard_logging_object = mocker.spy(
169+
litellm.litellm_core_utils.litellm_logging, "get_standard_logging_object_payload"
170+
)
171+
# First call (miss) - this will call the helper function
172+
await litellm.acompletion(model="gpt-3.5-turbo", messages=test_message, mock_response="r2", caching=True)
173+
await asyncio.sleep(0.1)
174+
mock_helper_fn.reset_mock()
175+
mock_get_standard_logging_object.reset_mock()
176+
177+
# Second call (hit) - this is the call we are testing
178+
await litellm.acompletion(model="gpt-3.5-turbo", messages=test_message, mock_response="r2", caching=True)
179+
await asyncio.sleep(0.1) # allow logs to process
180+
181+
# Assert the helper function was called exactly once during the cache hit
182+
mock_helper_fn.assert_called_once()
183+
assert mock_helper_fn.call_args.kwargs["cache_hit"] is True
184+
assert mock_helper_fn.call_args.kwargs["result"].choices[0].message.content == "r2"
185+
186+
mock_get_standard_logging_object.assert_called_once()
187+
assert mock_get_standard_logging_object.call_args.kwargs["status"] == "success"
188+
assert mock_get_standard_logging_object.call_args.kwargs["init_response_obj"].choices[0].message.content == "r2"

0 commit comments

Comments
 (0)