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

@@ -9,7 +9,7 @@
     "limit": 1078
   },
   "TQ004": {
-    "limit": 557
+    "limit": 544
   },
   "TQ005": {
     "limit": 2810

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

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