diff --git a/tests/logging_callback_tests/conftest.py b/tests/logging_callback_tests/conftest.py index 66d0ee01f8e..066afdf5c15 100644 --- a/tests/logging_callback_tests/conftest.py +++ b/tests/logging_callback_tests/conftest.py @@ -8,12 +8,18 @@ # globals like `litellm.num_retries = 3` which pollute state for all tests # in the same xdist worker. +import asyncio import importlib import os +from collections.abc import AsyncIterator +from typing import Final import pytest +import pytest_asyncio import litellm +from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, @@ -170,6 +176,15 @@ def isolate_litellm_state(): setattr(litellm, attr, _DEFAULTS[attr]) +LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0 + + +@pytest_asyncio.fixture(loop_scope="function", autouse=True) +async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]: + yield + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + + @pytest.fixture(scope="module", autouse=True) def setup_and_teardown(): """ diff --git a/tests/logging_callback_tests/logging_worker_drain_canary.py b/tests/logging_callback_tests/logging_worker_drain_canary.py new file mode 100644 index 00000000000..bff29129d7d --- /dev/null +++ b/tests/logging_callback_tests/logging_worker_drain_canary.py @@ -0,0 +1,23 @@ +import asyncio +import queue +from typing import Final + +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + +RUNS: Final[queue.SimpleQueue[tuple[asyncio.AbstractEventLoop, asyncio.AbstractEventLoop]]] = queue.SimpleQueue() + + +async def record_run(queued_on: asyncio.AbstractEventLoop) -> None: + RUNS.put((queued_on, asyncio.get_running_loop())) + + +async def test_1_leaves_an_event_pending() -> None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(record_run(asyncio.get_running_loop())) + + +async def test_2_never_inherits_the_pending_event() -> None: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + queued_on, ran_on = RUNS.get_nowait() + assert RUNS.empty() + assert ran_on is queued_on + assert ran_on is not asyncio.get_running_loop() diff --git a/tests/logging_callback_tests/test_logging_worker_drain.py b/tests/logging_callback_tests/test_logging_worker_drain.py new file mode 100644 index 00000000000..e6fef9880a1 --- /dev/null +++ b/tests/logging_callback_tests/test_logging_worker_drain.py @@ -0,0 +1,17 @@ +import os +from pathlib import Path +from typing import Final + +from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + +CANARY_MODULE: Final = Path(__file__).with_name("logging_worker_drain_canary.py") +CANARY_RUN: Final = ( + "import pytest\n" + f"raise SystemExit(pytest.main([{str(CANARY_MODULE)!r}, '-p', 'no:xdist', '-p', 'no:cacheprovider', '-q']))\n" +) + + +def test_drain_fixture_runs_pending_events_before_the_next_test_starts() -> None: + env_without_xdist: Final = {key: value for key, value in os.environ.items() if not key.startswith("PYTEST_XDIST")} + result: Final = run_child_interpreter(CANARY_RUN, env=env_without_xdist, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr