fix: complete missing testing

Author: justinmerrell

runpod/user_agent.py

@@ -7,7 +7,11 @@
 
 
 def construct_user_agent():
-    """ Constructs the User-Agent string for the RunPod-Python-SDK """
+    """ Constructs the User-Agent string for the RunPod-Python-SDK
+
+    Example:
+        RunPod-Python-SDK/0.1.0 (Linux 5.4.0-54-generic; x86_64) Language/Python 3.8.5
+    """
     os_info = f"{platform.system()} {platform.release()}; {platform.machine()}"
     python_version = platform.python_version()
     integration_method = os.getenv('RUNPOD_UA_INTEGRATION')

tests/test_serverless/test_worker.py

@@ -39,11 +39,10 @@ def test_get_auth_header(self):
         os_info = f"{platform.system()} {platform.release()}; {platform.machine()}"
         with patch("runpod.serverless.worker.os") as mock_os:
             mock_os.environ.get.return_value = "test"
-            mock_os.getenv.return_value = "test"
             assert runpod.serverless.worker._get_auth_header(
             ) == {
                 'Authorization': 'test',
-                'User-Agent': f'RunPod-Python-SDK/{runpod_version} ({os_info}) Language/Python {platform.python_version()} Integration/test'  # pylint: disable=line-too-long
+                'User-Agent': f'RunPod-Python-SDK/{runpod_version} ({os_info}) Language/Python {platform.python_version()}'  # pylint: disable=line-too-long
             }
 
     def test_is_local(self):

tests/test_user_agent.py

@@ -0,0 +1,49 @@
+""" Tests for the user_agent module. """
+
+import unittest
+from unittest.mock import patch
+import os
+
+from runpod.user_agent import construct_user_agent
+
+
+class TestConstructUserAgent(unittest.TestCase):
+    """Test the construct_user_agent function."""
+
+    @patch('runpod.user_agent.platform.system', return_value='Windows')
+    @patch('runpod.user_agent.platform.release', return_value='10')
+    @patch('runpod.user_agent.platform.machine', return_value='AMD64')
+    @patch('runpod.user_agent.platform.python_version', return_value='3.8.10')
+    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']
+
+        expected_ua = "RunPod-Python-SDK/0.1.0 (Windows 10; AMD64) Language/Python 3.8.10"
+        self.assertEqual(construct_user_agent(), expected_ua)
+
+        assert mock_python_version.called
+        assert mock_machine.called
+        assert mock_release.called
+        assert mock_system.called
+
+    @patch('runpod.user_agent.platform.system', return_value='Linux')
+    @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 = "RunPod-Python-SDK/0.1.0 (Linux 5.4; x86_64) Language/Python 3.9.5 Integration/SkyPilot"  # pylint: disable=line-too-long
+        self.assertEqual(construct_user_agent(), expected_ua)
+
+        assert mock_python_version.called
+        assert mock_machine.called
+        assert mock_release.called
+        assert mock_system.called
+
+
+if __name__ == '__main__':
+    unittest.main()