scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test(logging): cover async_failure_handler's correlation-context restore

Browse source 
Codecov flagged the new async_failure_handler wrapper (try/finally around
_async_failure_handler_body) as uncovered - the method had no direct test
at all before this PR's refactor split it into a wrapper. Adds a test that
awaits it directly and asserts both that async_log_failure_event still
fires and that _restore_correlation_context() puts the pre-call
trace_id/session_id back.
This commit is contained in:
Deepanshu 2026-07-23 17:41:31 -04:00
parent 313e170f1f
commit cc40bdb609
1 changed files with 41 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -3018,6 +3018,47 @@ def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(
dummy_logger.log_failure_event.assert_called_once()
@pytest.mark.asyncio
async def test_async_failure_handler_runs_callbacks_and_restores_correlation_context(logging_obj):
"""await logging_obj.async_failure_handler(...) must dispatch async failure callbacks
and, once its own body completes, restore trace_id/session_id contextvars via
_restore_correlation_context() (the fix for the nested-call context leak)."""
from litellm._logging import session_id_var, trace_id_var
from litellm.integrations.custom_logger import CustomLogger
class DummyLogger(CustomLogger):
pass
logging_obj.call_type = "acompletion"
logging_obj.stream = False
logging_obj.model_call_details["litellm_params"] = {}
logging_obj.litellm_params = {}
dummy_logger = DummyLogger()
dummy_logger.async_log_failure_event = AsyncMock()
trace_id_var.set("pre-existing-trace")
session_id_var.set("pre-existing-session")
try:
with patch.object(
logging_obj,
"get_combined_callback_list",
return_value=[dummy_logger],
):
await logging_obj.async_failure_handler(
exception=Exception("test error"),
traceback_exception="",
)
dummy_logger.async_log_failure_event.assert_called_once()
assert logging_obj._correlation_context_restored is True
assert trace_id_var.get() == "pre-existing-trace"
assert session_id_var.get() == "pre-existing-session"
finally:
trace_id_var.set("")
session_id_var.set("")
def test_merge_hidden_params_from_response_into_metadata_populates_metadata():
"""Streaming completion path should mirror non-stream: metadata.hidden_params from response."""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj