Merge pull request #38274 from BerriAI/litellm_test_lint_os_environ

test: gate the test tree on B003 so a test cannot swap os.environ for a plain dict

Author: ryan-crabbe-berri

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",
 ]

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
-    )
-
-    assert set(valid_models) == set(expected_models)
+        # list of openai supported llms on litellm
+        expected_models = (
+            litellm.open_ai_chat_completion_models | litellm.open_ai_text_completion_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
-
-    valid_models = get_valid_models()
-
-    print(valid_models)
-    assert set(valid_models) == set(expected_models)
+    with mock.patch.dict(os.environ, {"GEMINI_API_KEY": "temp"}, clear=True):
+        valid_models = get_valid_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"])