fix(tests): use threading.Event in MetadataCaptureCallback to avoid xdist cross-loop failures

asyncio.Event is bound to the event loop it was created in. Under pytest-xdist
with -n 16, each test function gets a fresh event loop but the LoggingWorker
may dispatch set() from a different loop context, causing the event to never
signal. threading.Event.set() is loop-agnostic and works reliably across all
parallel worker configurations. Await it via run_in_executor so the async test
can wait without blocking.
This commit is contained in:
Ishaan Jaffer 2026-03-07 17:34:58 -08:00
parent 5700c06689
commit 35cc601d4e

View file

@ -12,6 +12,7 @@ verifies metadata is preserved for custom callbacks via kwargs['litellm_params']
import asyncio
import os
import sys
import threading
from typing import Optional
from unittest.mock import AsyncMock, patch
@ -44,13 +45,16 @@ class MetadataCaptureCallback(CustomLogger):
def __init__(self):
self.captured_kwargs: Optional[dict] = None
self.event = asyncio.Event()
# Use threading.Event so set() works regardless of which event loop
# (or no loop) the logging worker uses — asyncio.Event is loop-bound
# and fails under pytest-xdist -n >1 where each test gets a fresh loop.
self._event = threading.Event()
async def async_log_success_event(
self, kwargs, response_obj, start_time, end_time
):
self.captured_kwargs = kwargs
self.event.set()
self._event.set()
@pytest.mark.asyncio
@ -106,7 +110,12 @@ async def test_metadata_passed_to_custom_callback_codex_models():
metadata=test_metadata,
)
await asyncio.wait_for(callback.event.wait(), timeout=5.0)
loop = asyncio.get_event_loop()
received = await asyncio.wait_for(
loop.run_in_executor(None, lambda: callback._event.wait(timeout=5.0)),
timeout=6.0,
)
assert received, "Callback was not invoked within timeout"
assert callback.captured_kwargs is not None, "Callback should have been invoked"
@ -167,7 +176,12 @@ async def test_metadata_passed_via_litellm_metadata_responses_api():
litellm_metadata=test_metadata,
)
await asyncio.wait_for(callback.event.wait(), timeout=5.0)
loop = asyncio.get_event_loop()
received = await asyncio.wait_for(
loop.run_in_executor(None, lambda: callback._event.wait(timeout=5.0)),
timeout=6.0,
)
assert received, "Callback was not invoked within timeout"
assert callback.captured_kwargs is not None