fix: refactor

justinmerrell · justinmerrell · commit 3fefbb57947c · 2024-02-12T21:14:33.000-05:00
diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py
@@ -184,18 +184,23 @@ async def run_job_generator(
     Run generator job used to stream output.
     Yields output partials from the generator.
     '''
+    is_async_gen = inspect.isasyncgenfunction(handler)
+    log.debug('Async Generator' if is_async_gen else 'Standard Generator', job["id"])
+
     try:
         job_output = handler(job)
-        if inspect.isasyncgenfunction(handler):
-            log.debug('Async generator', job["id"])
+
+        if is_async_gen:
             async for output_partial in job_output:
                 yield {"output": output_partial}
         else:
-            log.debug('Generator', job["id"])
             for output_partial in job_output:
                 yield {"output": output_partial}
+
     except Exception as err:    # pylint: disable=broad-except
         log.error(err, job["id"])
-        yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"}
+        yield {
+            "error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"
+        }
     finally:
         log.info('Finished running generator.', job["id"])
diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py
@@ -164,6 +164,23 @@ def generator_handler_exception(job):
     raise Exception()  # pylint: disable=broad-exception-raised
 
 
+def test_generator_handler_exception():
+    """ Test generator_handler_exception """
+    job = {"id": "test_job"}
+    gen = generator_handler_exception(job)
+
+    # Process the first yielded value
+    output = next(gen)
+    assert output == "test1", "First output should be 'test1'"
+
+    # Attempt to get the next value, expecting an exception
+    try:
+        next(gen)
+        assert False, "Expected an exception to be raised"
+    except Exception:  # pylint: disable=broad-except
+        assert True, "Exception was caught as expected"
+
+
 class TestRunWorker(IsolatedAsyncioTestCase):
     """ Tests for runpod | serverless| worker """
 
