From 3ac339cfbbe199e1667942cf082f7025505fe293 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng@berri.ai>
Date: Fri, 21 Aug 2026 22:10:34 -0700
Subject: [PATCH] test: stop the zai tests from leaking env and litellm globals
 into the session (#37834)

test_zai_provider.py set LITELLM_LOCAL_MODEL_COST_MAP and litellm.model_cost
directly and never put them back, so every test that ran after it in the same
process saw a local cost map instead of the real one. The two respx tests did
the same to litellm.disable_aiohttp_transport with no restore at all.

Both now go through monkeypatch, which restores on teardown including when the
test fails. The cost-map setup moves into a fixture requested by exactly the
five tests that read the cost map.
---
 test-quality-budget.json                      |  2 +-
 .../llms/zai/test_zai_provider.py             | 40 ++++++-------------
 2 files changed, 14 insertions(+), 28 deletions(-)

diff --git a/test-quality-budget.json b/test-quality-budget.json
index 6a62d783044..fd984fd1e21 100644
--- a/test-quality-budget.json
+++ b/test-quality-budget.json
@@ -12,7 +12,7 @@
     "limit": 469
   },
   "TQ005": {
-    "limit": 2549
+    "limit": 2542
   },
   "TQ006": {
     "limit": 34
diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py
index 61e1121257c..38ddac8d510 100644
--- a/tests/test_litellm/llms/zai/test_zai_provider.py
+++ b/tests/test_litellm/llms/zai/test_zai_provider.py
@@ -13,6 +13,12 @@ from litellm import completion
 from litellm.cost_calculator import cost_per_token
 
 
+@pytest.fixture
+def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
+
+
 @pytest.fixture
 def zai_response():
     """Mock response from Z.AI API"""
@@ -51,12 +57,8 @@ def test_zai_in_provider_lists():
     assert "zai" in litellm.provider_list
 
 
-def test_zai_models_in_model_cost(monkeypatch):
+def test_zai_models_in_model_cost(local_model_cost_map):
     """Test that ZAI models are in the model cost map"""
-    import os
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
 
     zai_models = [
         "zai/glm-4.7",
@@ -75,12 +77,8 @@ def test_zai_models_in_model_cost(monkeypatch):
         assert litellm.model_cost[model]["litellm_provider"] == "zai"
 
 
-def test_zai_glm46_cost_calculation(monkeypatch):
+def test_zai_glm46_cost_calculation(local_model_cost_map):
     """Test the cost calculation for glm-4.6"""
-    import os
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
 
     key = "zai/glm-4.6"
     info = litellm.model_cost[key]
@@ -96,12 +94,8 @@ def test_zai_glm46_cost_calculation(monkeypatch):
     assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
 
 
-def test_zai_flash_model_is_free(monkeypatch):
+def test_zai_flash_model_is_free(local_model_cost_map):
     """Test that glm-4.5-flash has zero cost"""
-    import os
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
 
     key = "zai/glm-4.5-flash"
     info = litellm.model_cost[key]
@@ -110,12 +104,8 @@ def test_zai_flash_model_is_free(monkeypatch):
     assert info["output_cost_per_token"] == 0
 
 
-def test_glm47_supports_reasoning(monkeypatch):
+def test_glm47_supports_reasoning(local_model_cost_map):
     """Test that GLM-4.7 supports reasoning"""
-    import os
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
 
     key = "zai/glm-4.7"
     assert key in litellm.model_cost, f"Model {key} not found in model_cost"
@@ -124,12 +114,8 @@ def test_glm47_supports_reasoning(monkeypatch):
     assert info["supports_reasoning"] is True
 
 
-def test_glm47_cost_calculation(monkeypatch):
+def test_glm47_cost_calculation(local_model_cost_map):
     """Test cost calculation for GLM-4.7"""
-    import os
-
-    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-    litellm.model_cost = litellm.get_model_cost_map(url="")
 
     prompt_cost, completion_cost = cost_per_token(
         model="zai/glm-4.7",
@@ -146,7 +132,7 @@ def test_glm47_cost_calculation(monkeypatch):
 async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
     """Test completion call with zai provider using mocked response"""
     monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
-    litellm.disable_aiohttp_transport = True
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
 
     respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
         json=zai_response
@@ -172,7 +158,7 @@ async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
 def test_zai_sync_completion(respx_mock, zai_response, monkeypatch):
     """Test synchronous completion call"""
     monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
-    litellm.disable_aiohttp_transport = True
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
 
     respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
         json=zai_response