|
30 | 30 | )
|
31 | 31 | from litellm.integrations.custom_logger import CustomLogger
|
32 | 32 | from litellm.types.utils import StandardLoggingPayload
|
| 33 | +from litellm.litellm_core_utils.litellm_logging import Logging |
33 | 34 |
|
34 | 35 |
|
35 | 36 | class CacheHitCustomLogger(CustomLogger):
|
@@ -152,3 +153,36 @@ async def test_cache_hit_includes_custom_llm_provider():
|
152 | 153 | # Clean up
|
153 | 154 | litellm.callbacks = original_callbacks
|
154 | 155 | 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