Update test_worker.py

Author: justinmerrell

tests/test_serverless/test_worker.py

@@ -208,7 +208,8 @@ def test_run_worker(
         mock_run_job.return_value = {"output": {"result": "odd"}}
 
         # Call the function
-        runpod.serverless.start(self.config)
+        with patch("runpod.serverless.os", return_value=None):
+            runpod.serverless.start(self.config)
 
         # Make assertions about the behaviors
         mock_get_job.assert_called_once()
@@ -241,7 +242,8 @@ def test_run_worker_generator_handler(
         # Test generator handler
         generator_config = {
             "handler": generator_handler, "refresh_worker": True}
-        runpod.serverless.start(generator_config)
+        with patch("runpod.serverless.os", return_value=None):
+            runpod.serverless.start(generator_config)
 
         assert mock_stream_result.called
         assert not mock_run_job.called
@@ -273,7 +275,8 @@ def test_run_worker_generator_handler_exception(
         # Test generator handler
         generator_config = {
             "handler": generator_handler_exception, "refresh_worker": True}
-        runpod.serverless.start(generator_config)
+        with patch("runpod.serverless.os", return_value=None):
+            runpod.serverless.start(generator_config)
 
         assert mock_stream_result.call_count == 1
         assert not mock_run_job.called
@@ -306,7 +309,8 @@ def test_run_worker_generator_aggregate_handler(
         # Test generator handler
         generator_config = {
             "handler": generator_handler, "return_aggregate_stream": True, "refresh_worker": True}
-        runpod.serverless.start(generator_config)
+        with patch("runpod.serverless.os", return_value=None):
+            runpod.serverless.start(generator_config)
 
         assert mock_send_result.called
         assert mock_stream_result.called
@@ -345,7 +349,8 @@ def concurrency_modifier(current_concurrency):
         config_with_concurrency['concurrency_modifier'] = concurrency_modifier
 
         # Call the function
-        runpod.serverless.start(config_with_concurrency)
+        with patch("runpod.serverless.os", return_value=None):
+            runpod.serverless.start(config_with_concurrency)
 
         # Make assertions about the behaviors
         mock_get_job.assert_called_once()
@@ -379,7 +384,8 @@ def test_run_worker_multi_processing(
         mock_run_job.return_value = {"output": {"result": "odd"}}
 
         # Call the function
-        runpod.serverless.start(self.config)
+        with patch("runpod.serverless.os", return_value=None):
+            runpod.serverless.start(self.config)
 
         # Make assertions about the behaviors
         mock_get_job.assert_called_once()