From 74816498303d18451d7452fbcabe3042e1ba7de2 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng@berri.ai>
Date: Fri, 21 Aug 2026 20:39:07 -0700
Subject: [PATCH] test(datadog): restore an empty DD_API_KEY instead of
 unsetting it (#37832)

Both datadog test files hand-roll what monkeypatch.setenv already does: read the
old value, write the test value, put the old one back on the way out. The cost
management fixture checks the old value for truthiness rather than for None, so
an operator running the suite with DD_API_KEY set to the empty string gets it
deleted rather than restored. Starting from DD_API_KEY="" and running test_init
leaves it None on the current file, and "" after this.

13 raw os.environ writes become monkeypatch.setenv, the two fixtures stop being
yield fixtures because there is nothing left to do on the way out, and the now
unused os import goes with them.

27 tests pass across the two files, 88 across tests/test_litellm/integrations/datadog.
---
 test-quality-budget.json                      |  2 +-
 .../datadog/test_datadog_cost_management.py   | 36 ++++---------------
 .../datadog/test_datadog_metrics.py           | 30 ++++++----------
 3 files changed, 18 insertions(+), 50 deletions(-)

diff --git a/test-quality-budget.json b/test-quality-budget.json
index 143378efec7..68244dce319 100644
--- a/test-quality-budget.json
+++ b/test-quality-budget.json
@@ -9,7 +9,7 @@
     "limit": 1078
   },
   "TQ004": {
-    "limit": 557
+    "limit": 544
   },
   "TQ005": {
     "limit": 2810
diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py
index cb786d9c292..1a50a6991da 100644
--- a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py
+++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py
@@ -1,4 +1,3 @@
-import os
 import time
 from unittest.mock import AsyncMock
 
@@ -12,34 +11,13 @@ from litellm.types.utils import StandardLoggingPayload
 
 
 @pytest.fixture
-def clean_env():
-    # Save original env
-    original_api_key = os.environ.get("DD_API_KEY")
-    original_app_key = os.environ.get("DD_APP_KEY")
-    original_site = os.environ.get("DD_SITE")
-
-    # Set test env
-    os.environ["DD_API_KEY"] = "test_api_key"
-    os.environ["DD_APP_KEY"] = "test_app_key"
-    os.environ["DD_SITE"] = "test.datadoghq.com"
-
-    yield
-
-    # Restore original env
-    if original_api_key:
-        os.environ["DD_API_KEY"] = original_api_key
-    else:
-        del os.environ["DD_API_KEY"]
-
-    if original_app_key:
-        os.environ["DD_APP_KEY"] = original_app_key
-    else:
-        del os.environ["DD_APP_KEY"]
-
-    if original_site:
-        os.environ["DD_SITE"] = original_site
-    else:
-        del os.environ["DD_SITE"]
+def clean_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    for key, value in (
+        ("DD_API_KEY", "test_api_key"),
+        ("DD_APP_KEY", "test_app_key"),
+        ("DD_SITE", "test.datadoghq.com"),
+    ):
+        monkeypatch.setenv(key, value)
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py
index a4a4ca334b0..eade92d6672 100644
--- a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py
+++ b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py
@@ -1,4 +1,3 @@
-import os
 import time
 from datetime import datetime, timedelta
 from unittest.mock import AsyncMock
@@ -11,25 +10,16 @@ from litellm.types.utils import StandardLoggingPayload
 
 
 @pytest.fixture
-def clean_env():
-    """Set test env vars and restore originals after test."""
-    keys = ["DD_API_KEY", "DD_APP_KEY", "DD_SITE", "DD_ENV", "DD_SERVICE", "DD_VERSION"]
-    originals = {k: os.environ.get(k) for k in keys}
-
-    os.environ["DD_API_KEY"] = "test_api_key"
-    os.environ["DD_APP_KEY"] = "test_app_key"
-    os.environ["DD_SITE"] = "test.datadoghq.com"
-    os.environ["DD_ENV"] = "test-env"
-    os.environ["DD_SERVICE"] = "test-service"
-    os.environ["DD_VERSION"] = "1.0.0"
-
-    yield
-
-    for k, v in originals.items():
-        if v is not None:
-            os.environ[k] = v
-        elif k in os.environ:
-            del os.environ[k]
+def clean_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    for key, value in (
+        ("DD_API_KEY", "test_api_key"),
+        ("DD_APP_KEY", "test_app_key"),
+        ("DD_SITE", "test.datadoghq.com"),
+        ("DD_ENV", "test-env"),
+        ("DD_SERVICE", "test-service"),
+        ("DD_VERSION", "1.0.0"),
+    ):
+        monkeypatch.setenv(key, value)
 
 
 @pytest.mark.asyncio