diff --git a/ruff-tests.toml b/ruff-tests.toml index 7d28d6d838a..e036fb4946c 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -51,6 +51,10 @@ # F632 `is` against a literal. It compares identity, so it passes only where CPython # happens to intern the value and stops meaning what it says the moment the # value is built at runtime +# B003 `os.environ = {...}` rebinds the mapping instead of mutating it, so `putenv` +# never fires and a subprocess still reads the real keys the test believes it +# cleared. The manual restore underneath is skipped whenever the body raises, +# so every later test in that worker inherits a plain dict for an environment # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -78,4 +82,5 @@ lint.select = [ "B023", "B025", "F632", + "B003", ] diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 67f2e1ce06d..0ccfae55290 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -328,34 +328,23 @@ def test_trimming_with_untokenizable_field(caplog: pytest.LogCaptureFixture) -> def test_aget_valid_models(): - old_environ = os.environ - os.environ = {"OPENAI_API_KEY": "temp"} # mock set only openai key in environ + with mock.patch.dict(os.environ, {"OPENAI_API_KEY": "temp"}, clear=True): + valid_models = get_valid_models() + print(valid_models) - valid_models = get_valid_models() - print(valid_models) + # list of openai supported llms on litellm + expected_models = ( + litellm.open_ai_chat_completion_models | litellm.open_ai_text_completion_models + ) - # list of openai supported llms on litellm - expected_models = ( - litellm.open_ai_chat_completion_models | litellm.open_ai_text_completion_models - ) - - assert set(valid_models) == set(expected_models) - - # reset replicate env key - os.environ = old_environ + assert set(valid_models) == set(expected_models) # GEMINI - expected_models = litellm.gemini_models - old_environ = os.environ - os.environ = {"GEMINI_API_KEY": "temp"} # mock set only openai key in environ + with mock.patch.dict(os.environ, {"GEMINI_API_KEY": "temp"}, clear=True): + valid_models = get_valid_models() - valid_models = get_valid_models() - - print(valid_models) - assert set(valid_models) == set(expected_models) - - # reset replicate env key - os.environ = old_environ + print(valid_models) + assert set(valid_models) == set(litellm.gemini_models) @pytest.mark.parametrize("custom_llm_provider", ["anthropic", "xai"])