fix: tests

Author: justinmerrell

tests/test_serverless/test_modules/test_logger.py

@@ -127,9 +127,9 @@ def test_log_job_id(self):
             # Test with endpoint id set
             os.environ["RUNPOD_ENDPOINT_ID"] = "test_endpoint_id"
             logger.log("test_message", "INFO", job_id)
-            os.environ.pop("RUNPOD_ENDPOINT_ID")
 
-            mock_print.assert_called_once_with(
+            mock_print.assert_called_with(
                 '{"requestId": "test_job_id", "message": "test_message", "level": "INFO"}',
                 flush=True
             )
+            os.environ.pop("RUNPOD_ENDPOINT_ID")

tests/test_serverless/test_worker.py

@@ -4,6 +4,7 @@
 import os
 import platform
 import argparse
+from unittest import mock
 from unittest.mock import patch, mock_open, Mock, MagicMock
 
 from unittest import IsolatedAsyncioTestCase
@@ -18,24 +19,17 @@
 
 
 class TestWorker(IsolatedAsyncioTestCase):
-    """ Tests for runpod | serverless| worker """
+    """Tests for RunPod serverless worker."""
 
     def setUp(self):
-        self.mock_handler = Mock()
-        self.mock_handler.return_value = "test"
-
-        self.mock_config = Mock()
+        self.mock_handler = mock.Mock(return_value="test")
         self.mock_config = {
             "handler": self.mock_handler,
-            "rp_args": {
-                "test_input": None,
-            }
+            "rp_args": {"test_input": None},
         }
 
     def test_get_auth_header(self):
-        '''
-        Test _get_auth_header
-        '''
+        """Test retrieval of the authentication header from _get_auth_header."""
         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"