diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 9028b8461b0..c96e3e2c031 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -307,6 +307,30 @@ def isolate_router_model_cost_state() -> Iterator[None]: litellm.get_model_info.cache_clear() +def _provider_model_sets() -> tuple[set[str], ...]: + module_globals: Final = tuple(vars(litellm).items()) + return tuple(models for name, models in module_globals if name.endswith("_models") and isinstance(models, set)) + + +@pytest.fixture(autouse=True) +def isolate_model_catalog() -> Iterator[None]: + model_cost: Final = litellm.model_cost + model_cost_entries: Final = dict(model_cost) + models_by_provider: Final = litellm.models_by_provider + models_by_provider_entries: Final = dict(models_by_provider) + provider_model_sets: Final = tuple((models, frozenset(models)) for models in _provider_model_sets()) + yield + model_cost.clear() + model_cost.update(model_cost_entries) + models_by_provider.clear() + models_by_provider.update(models_by_provider_entries) + for models, entries in provider_model_sets: + models.clear() + models.update(entries) + 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") diff --git a/tests/unit/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py index 63db0eb5bba..92457119a82 100644 --- a/tests/unit/test_conftest_isolation.py +++ b/tests/unit/test_conftest_isolation.py @@ -88,6 +88,26 @@ def test_register_model_ledger_entry_was_rolled_back(): assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost +def test_register_model_catalog_writes_are_scoped_to_this_test(): + litellm.register_model( + { + CANARY_MODEL: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.001}, + "collection-isolation-ledger": {"input_cost_per_token": 0.5}, + } + ) + assert CANARY_MODEL in litellm.model_cost + assert CANARY_MODEL in litellm.open_ai_chat_completion_models + assert CANARY_MODEL in litellm.models_by_provider["openai"] + assert litellm.model_cost["collection-isolation-ledger"]["input_cost_per_token"] == 0.5 + + +def test_register_model_catalog_writes_were_rolled_back(): + assert CANARY_MODEL not in litellm.model_cost + assert CANARY_MODEL not in litellm.open_ai_chat_completion_models + assert CANARY_MODEL not in litellm.models_by_provider["openai"] + assert litellm.model_cost["collection-isolation-ledger"]["input_cost_per_token"] == 0.001 + + def test_live_router_membership_is_scoped_to_this_test(): _CanaryRouterHolder.router = Router( model_list=[ diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0be3ae05e97..40eb5c31fda 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2358,7 +2358,6 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() -@pytest.mark.usefixtures("isolated_openai_model_sets") @pytest.mark.usefixtures("isolated_openai_model_sets") @respx.mock def test_register_model_url_fetch_uses_single_attempt(monkeypatch): @@ -4803,8 +4802,7 @@ async def test_builtin_string_callback_registers_when_subclass_already_active( assert any(type(cb) is S3Logger for cb in litellm._async_success_callback) -@pytest.mark.usefixtures("isolated_openai_model_sets") -@pytest.mark.usefixtures("isolated_openai_model_sets") +@pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") @pytest.mark.parametrize("rate", (0.0, 0.000123)) @pytest.mark.parametrize("updated_rate", (None, 0.0, 0.000456)) def test_reapply_runtime_registrations_replays_register_model_overrides( @@ -4871,8 +4869,7 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( _invalidate_model_cost_lowercase_map() -@pytest.mark.usefixtures("isolated_openai_model_sets") -@pytest.mark.usefixtures("isolated_openai_model_sets") +@pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkeypatch): """ Per-request custom pricing describes one call, so it must not be re-asserted