diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index f4aa1926d21..c0993051d33 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -19,6 +19,7 @@ sys.path.insert( import asyncio import litellm +from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( image_handling as image_handling_module, @@ -238,6 +239,11 @@ def isolate_litellm_state(): if hasattr(litellm, _attr): original_state[_attr] = getattr(litellm, _attr) + original_runtime_registered_model_cost = { + model_key: dict(model_value) + for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() + } + # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for # JSON logging and do not restore them, which breaks later caplog-based tests. logger_state = {} @@ -304,6 +310,9 @@ def isolate_litellm_state(): if hasattr(litellm, attr_name): setattr(litellm, attr_name, original_value) + litellm_utils_module._runtime_registered_model_cost.clear() + litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/test_litellm/test_conftest_isolation.py new file mode 100644 index 00000000000..88889ad7740 --- /dev/null +++ b/tests/test_litellm/test_conftest_isolation.py @@ -0,0 +1,13 @@ +import litellm +from litellm import utils as litellm_utils_module + +CANARY_MODEL = "conftest-isolation-canary-model" + + +def test_register_model_ledger_entry_is_scoped_to_this_test(): + litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) + assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost + + +def test_register_model_ledger_entry_was_rolled_back(): + assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost