Merge pull request #81 from runpod/fix_job_gen

fix this

Author: justinmerrell

runpod/serverless/modules/rp_job.py

@@ -155,8 +155,12 @@ async def run_job_generator(
     '''
     try:
         job_output = handler(job)
-        for output_partial in job_output:
-            yield {"output": output_partial}
+        if inspect.isasyncgenfunction(handler):
+            async for output_partial in job_output:
+                yield {"output": output_partial}
+        else:
+            for output_partial in job_output:
+                yield {"output": output_partial}
     except Exception as err:    # pylint: disable=broad-except
         log.error(f'Error while running job {job["id"]}: {err}')
         yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"}

tests/test_serverless/test_modules/test_job.py

@@ -214,24 +214,46 @@ async def test_job_with_exception(self):
 class TestRunJobGenerator(IsolatedAsyncioTestCase):
     ''' Tests the run_job_generator function '''
 
-    def handler_success(self, job): # pylint: disable=unused-argument
+    def handler_gen_success(self, job): # pylint: disable=unused-argument
         '''
-        Test handler that returns a generator
+        Test handler that returns a generator.
+        '''
+        yield "partial_output_1"
+        yield "partial_output_2"
+
+    async def handler_async_gen_success(self, job): # pylint: disable=unused-argument
+        '''
+        Test handler that returns an async generator.
         '''
         yield "partial_output_1"
         yield "partial_output_2"
 
     def handler_fail(self, job):
         '''
-        Test handler that raises an exception
+        Test handler that raises an exception.
         '''
         raise Exception("Test Exception") # pylint: disable=broad-exception-raised
 
     async def test_run_job_generator_success(self):
         '''
         Tests the run_job_generator function with a successful generator
         '''
-        handler = self.handler_success
+        handler = self.handler_gen_success
+        job = {"id": "123"}
+
+        with patch("runpod.serverless.modules.rp_job.log", new_callable=Mock) as mock_log:
+            result = [i async for i in rp_job.run_job_generator(handler, job)]
+
+        assert result == [{"output": "partial_output_1"}, {"output": "partial_output_2"}]
+        assert mock_log.error.call_count == 0
+        assert mock_log.info.call_count == 1
+        mock_log.info.assert_called_with('123 | Finished ')
+
+    async def test_run_job_generator_success_async(self):
+        '''
+        Tests the run_job_generator function with a successful generator
+        '''
+        handler = self.handler_async_gen_success
         job = {"id": "123"}
 
         with patch("runpod.serverless.modules.rp_job.log", new_callable=Mock) as mock_log: