mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
784a52de79
* fix(tests): keep foreign asyncio task noise out of the budget limiter caplog asserts The two gc.collect()-under-caplog tests in test_budget_limiter_hotpath.py asserted on every captured record, so a LoggingWorker task another test had parked on the session fixture loop got destroyed inside their window and asyncio's "Task was destroyed but it is pending!" error failed them. They now assert only the limiter logger's records and park a pending task on a closed loop in the window themselves, so the noise is exercised on every run The edenai and sail spend_capture fixtures no longer start the worker from the session loop: the test's own call and its flush bind the worker to the test loop * test(router_strategy): drop only asyncio records in the limiter caplog asserts and drain the response cache worker on the test loop --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
import asyncio
|
|
import uuid
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
|
|
|
|
|
@pytest.fixture
|
|
def eden_key(monkeypatch) -> str:
|
|
monkeypatch.delenv("EDENAI_API_BASE", raising=False)
|
|
monkeypatch.setenv("EDENAI_API_KEY", "eden-test-key")
|
|
monkeypatch.setattr(litellm, "api_key", None)
|
|
return "eden-test-key"
|
|
|
|
|
|
@pytest.fixture
|
|
def no_eden_key(monkeypatch) -> None:
|
|
monkeypatch.delenv("EDENAI_API_KEY", raising=False)
|
|
monkeypatch.setattr(litellm, "api_key", None)
|
|
|
|
|
|
class SpendCapture(CustomLogger):
|
|
"""Records the cost the spend logs would store for one call, matched by its call id."""
|
|
|
|
def __init__(self, call_id: str):
|
|
super().__init__()
|
|
self.call_id = call_id
|
|
self.costs: list[object] = []
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
if kwargs.get("litellm_call_id") == self.call_id:
|
|
self.costs.append((kwargs.get("standard_logging_object") or {}).get("response_cost"))
|
|
|
|
async def settle(self) -> None:
|
|
await asyncio.sleep(0)
|
|
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
|
|
|
|
|
|
@pytest.fixture
|
|
def spend_capture(monkeypatch) -> SpendCapture:
|
|
capture = SpendCapture(call_id=f"eden-{uuid.uuid4()}")
|
|
monkeypatch.setattr(litellm, "callbacks", [capture])
|
|
return capture
|
|
|
|
|
|
@pytest.fixture
|
|
def httpx_transport(monkeypatch):
|
|
"""respx fakes httpx, so the async client must not sit on LiteLLM's default aiohttp transport."""
|
|
monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary.
|
|
litellm,
|
|
"disable_aiohttp_transport",
|
|
True,
|
|
)
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|
|
yield
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|