diff --git a/tests/unit/llms/anthropic/messages/native/test_response_cache.py b/tests/unit/llms/anthropic/messages/native/test_response_cache.py index 133bcd96433..bc3b1f6ac76 100644 --- a/tests/unit/llms/anthropic/messages/native/test_response_cache.py +++ b/tests/unit/llms/anthropic/messages/native/test_response_cache.py @@ -2,6 +2,7 @@ import asyncio from typing import Any, AsyncIterator, Dict, Final, List import pytest +import pytest_asyncio import datetime @@ -67,7 +68,7 @@ async def _collect(stream: AsyncIterator[bytes]) -> List[bytes]: return [chunk async for chunk in stream] -@pytest.fixture(autouse=True) +@pytest_asyncio.fixture(autouse=True, loop_scope="function") async def _drain_logging_worker(): yield await GLOBAL_LOGGING_WORKER.flush() diff --git a/tests/unit/llms/edenai/conftest.py b/tests/unit/llms/edenai/conftest.py index 5ca5728354c..e11a515e953 100644 --- a/tests/unit/llms/edenai/conftest.py +++ b/tests/unit/llms/edenai/conftest.py @@ -2,7 +2,6 @@ import asyncio import uuid import pytest -import pytest_asyncio import litellm from litellm.integrations.custom_logger import CustomLogger @@ -40,9 +39,8 @@ class SpendCapture(CustomLogger): await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) -@pytest_asyncio.fixture -async def spend_capture(monkeypatch) -> SpendCapture: - GLOBAL_LOGGING_WORKER.start() # rebinds the worker's queue to this test's event loop +@pytest.fixture +def spend_capture(monkeypatch) -> SpendCapture: capture = SpendCapture(call_id=f"eden-{uuid.uuid4()}") monkeypatch.setattr(litellm, "callbacks", [capture]) return capture diff --git a/tests/unit/llms/sail/conftest.py b/tests/unit/llms/sail/conftest.py index 2b2a6e5cae0..c1156af5b12 100644 --- a/tests/unit/llms/sail/conftest.py +++ b/tests/unit/llms/sail/conftest.py @@ -4,11 +4,9 @@ from typing import Final import httpx import pytest -import pytest_asyncio import respx import litellm -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests.unit.llms.sail.helpers import SAIL_API_BASE, SpendCapture, chat_completion_body @@ -26,9 +24,8 @@ def sail_env(local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> Ite litellm.in_memory_llm_clients_cache.flush_cache() -@pytest_asyncio.fixture -async def spend_capture(monkeypatch: pytest.MonkeyPatch) -> SpendCapture: - GLOBAL_LOGGING_WORKER.start() +@pytest.fixture +def spend_capture(monkeypatch: pytest.MonkeyPatch) -> SpendCapture: capture: Final = SpendCapture(call_id=f"sail-{uuid.uuid4()}") monkeypatch.setattr(litellm, "callbacks", [capture]) return capture diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 10704a29347..c79ff6325d2 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -2,6 +2,7 @@ import asyncio import contextlib import gc import logging +import threading from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -301,6 +302,27 @@ def test_router_add_deployment_registers_deployment_budget(disable_budget_sync, assert config.max_budget == 0.000000000001 +def _leave_a_task_pending_on_a_closed_loop() -> None: + def park_a_task_then_close_the_loop() -> None: + loop: Final = asyncio.new_event_loop() + + async def park() -> None: + await loop.create_future() + + loop.create_task(park()) + loop.call_soon(loop.stop) + loop.run_forever() + loop.close() + + parker: Final = threading.Thread(target=park_a_task_then_close_the_loop) + parker.start() + parker.join() + + +def _non_asyncio_log_lines(caplog: pytest.LogCaptureFixture) -> list[str]: + return [record.getMessage() for record in caplog.records if record.name != "asyncio"] + + @pytest.mark.asyncio async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_task(disable_budget_sync, caplog): """The budget sync runs every second, so an open breaker must not add an error line or an unretrieved task exception per cycle.""" @@ -322,11 +344,12 @@ async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_ta with caplog.at_level(logging.ERROR): await limiter._sync_in_memory_spend_with_redis() await asyncio.sleep(0) + _leave_a_task_pending_on_a_closed_loop() gc.collect() finally: loop.set_exception_handler(None) - assert caplog.records == [] + assert _non_asyncio_log_lines(caplog) == [] unretrieved.assert_not_called() assert limiter.redis_increment_operation_queue == [ {"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60} @@ -382,11 +405,12 @@ async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sy with caplog.at_level(logging.ERROR): await limiter._push_in_memory_increments_to_redis() await asyncio.sleep(0) + _leave_a_task_pending_on_a_closed_loop() gc.collect() finally: loop.set_exception_handler(None) - assert [record.getMessage() for record in caplog.records if record.name != "asyncio"] == [ + assert _non_asyncio_log_lines(caplog) == [ "Error syncing in-memory cache with Redis: Error 61 connecting to 127.0.0.1:6379" ] unretrieved.assert_not_called()