From decb28f5b575378cff5e8ce941c999da24df3875 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:34:55 +0000 Subject: [PATCH] test(unit): restore live router and runtime model cost state between unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index d002452d14c..d47f71b21a8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -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")