fix: use explicit env var set/restore in user agent tests

Author: KAJdev

tests/test_user_agent.py

@@ -19,12 +19,13 @@ def test_user_agent_without_integration(
         self, mock_python_version, mock_machine, mock_release, mock_system
     ):
         """Test the User-Agent string without specifying an integration method."""
-        if "RUNPOD_UA_INTEGRATION" in os.environ:
-            del os.environ["RUNPOD_UA_INTEGRATION"]
+        saved = {k: os.environ.pop(k) for k in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE") if k in os.environ}
 
         expected_ua = f"RunPod-Python-SDK/{runpod_version} (Windows 10; AMD64) Language/Python 3.8.10"  # pylint: disable=line-too-long
         self.assertEqual(construct_user_agent(), expected_ua)
 
+        os.environ.update(saved)
+
         assert mock_python_version.called
         assert mock_machine.called
         assert mock_release.called
@@ -34,16 +35,19 @@ def test_user_agent_without_integration(
     @patch("runpod.user_agent.platform.release", return_value="5.4")
     @patch("runpod.user_agent.platform.machine", return_value="x86_64")
     @patch("runpod.user_agent.platform.python_version", return_value="3.9.5")
-    @patch.dict(os.environ, {"RUNPOD_UA_INTEGRATION": "SkyPilot"})
     def test_user_agent_with_integration(
         self, mock_python_version, mock_machine, mock_release, mock_system
     ):
         """Test the User-Agent string with an integration method specified."""
-        expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 Integration/SkyPilot"  # pylint: disable=line-too-long
-
+        saved_claude = os.environ.pop("CLAUDECODE", None)
         os.environ["RUNPOD_UA_INTEGRATION"] = "SkyPilot"
+
+        expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 Integration/SkyPilot"  # pylint: disable=line-too-long
         self.assertEqual(construct_user_agent(), expected_ua)
+
         os.environ.pop("RUNPOD_UA_INTEGRATION")
+        if saved_claude is not None:
+            os.environ["CLAUDECODE"] = saved_claude
 
         assert mock_python_version.called
         assert mock_machine.called
@@ -55,16 +59,16 @@ def test_user_agent_with_integration(
     @patch("runpod.user_agent.platform.release", return_value="5.4")
     @patch("runpod.user_agent.platform.machine", return_value="x86_64")
     @patch("runpod.user_agent.platform.python_version", return_value="3.9.5")
-    @patch.dict(os.environ, {"CLAUDECODE": "1"}, clear=False)
     def test_user_agent_with_claude_code(
         self, mock_python_version, mock_machine, mock_release, mock_system
     ):
         """Test the User-Agent string includes claude-code agent tag."""
-        if "RUNPOD_UA_INTEGRATION" in os.environ:
-            del os.environ["RUNPOD_UA_INTEGRATION"]
+        os.environ.pop("RUNPOD_UA_INTEGRATION", None)
+        os.environ["CLAUDECODE"] = "1"
 
         expected_ua = f"RunPod-Python-SDK/{runpod_version} (Linux 5.4; x86_64) Language/Python 3.9.5 (via claude-code)"
         self.assertEqual(construct_user_agent(), expected_ua)
+        os.environ.pop("CLAUDECODE", None)
 
     @patch("runpod.user_agent.platform.system", return_value="Linux")
     @patch("runpod.user_agent.platform.release", return_value="5.4")
@@ -74,13 +78,13 @@ def test_user_agent_without_claude_code(
         self, mock_python_version, mock_machine, mock_release, mock_system
     ):
         """Test the User-Agent string excludes agent tag when env var is not set."""
-        for key in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE"):
-            if key in os.environ:
-                del os.environ[key]
+        saved = {k: os.environ.pop(k) for k in ("RUNPOD_UA_INTEGRATION", "CLAUDECODE") if k in os.environ}
 
         ua = construct_user_agent()
         self.assertNotIn("via claude-code", ua)
 
+        os.environ.update(saved)
+
 
 if __name__ == "__main__":
     unittest.main()