scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

test(unit): restore live router and runtime model cost state between unit tests

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng 2026-09-20 08:34:55 +00:00
parent fcabb626ac
commit decb28f5b5
1 changed files with 21 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
tests/unit/conftest.py
View file

@ -2,9 +2,11 @@ from collections.abc import Iterator
from typing import Final
import pytest
from pytest_socket import enable_socket, socket_allow_hosts
import litellm
from pytest_socket import enable_socket, socket_allow_hosts
import litellm.router as litellm_router_module
import litellm.utils as litellm_utils_module
LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"]
@ -25,6 +27,24 @@ def pytest_runtest_setup() -> None:
_allow_loopback_only()
@pytest.fixture(autouse=True)
def isolate_router_model_cost_state() -> Iterator[None]:
original_live_routers: Final = frozenset(litellm_router_module._live_routers)
original_runtime_registered_model_cost: Final = {
model_key: dict(model_value)
for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items()
}
yield
for router in tuple(litellm_router_module._live_routers):
litellm_router_module._live_routers.discard(router)
for router in original_live_routers:
litellm_router_module._live_routers.add(router)
litellm_utils_module._runtime_registered_model_cost.clear()
litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost)
litellm_utils_module._invalidate_model_cost_lowercase_map()
litellm.get_model_info.cache_clear()
@pytest.fixture
def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")