finished async_test_batch_error_handling_multiple_runnables

Author: tomerm-iguazio

tests/test_flow.py

@@ -126,6 +126,21 @@ def run(self, data, path, origin_name=None):
         return random_uuid
 
 
+class RunnableRaiseIfNegative(ParallelExecutionRunnable):
+    def run(self, data, path, origin_name=None):
+        if isinstance(data, list):
+            results = []
+            for item in data:
+                if item < 0:
+                    raise ValueError(f"Value {item} is negative!")
+                results.append(item * 2)
+            return results
+        else:
+            if data < 0:
+                raise ValueError(f"Value {data} is negative!")
+            return data * 2
+
+
 class MyParallelExecution(ParallelExecution):
     def select_runnables(self, event):
         return ["multiply", "add", "uuid"]
@@ -2504,26 +2519,15 @@ async def emit_event(i):
             assert termination_result[i]["uuid"] == expected_uuid
         previous_batch_number = batch_number
 
+
 def test_batch_error_handling_single_runnable():
     asyncio.run(async_test_batch_error_handling_single_runnable())
 
+
 @pytest.mark.asyncio
 async def async_test_batch_error_handling_single_runnable():
     """Test error handling in batched parallel execution with single runnable."""
     flush_after_seconds = 0.3
-    class RunnableRaiseIfNegative(ParallelExecutionRunnable):
-        def run(self, data, path, origin_name=None):
-            if isinstance(data, list):
-                results = []
-                for item in data:
-                    if item < 0:
-                        raise ValueError(f"Value {item} is negative!")
-                    results.append(item * 2)
-                return results
-            else:
-                if data < 0:
-                    raise ValueError(f"Value {data} is negative!")
-                return data * 2
 
     batch_size = 3
     runnables = [RunnableRaiseIfNegative("check_positive", raise_exception=False)]
@@ -2535,26 +2539,29 @@ def run(self, data, path, origin_name=None):
         execution_mechanism_by_runnable_name={
             "check_positive": "naive",
         },
-
     )
     reducer = Reduce([], append_and_return)
 
-    controller = build_flow([
-        source,
-        batch_step,
-        parallel_execution,
-        FlatMap(fn=lambda x: x.body, full_event=True),
-        Complete(),
-        reducer,
-    ]).run()
+    controller = build_flow(
+        [
+            source,
+            batch_step,
+            parallel_execution,
+            FlatMap(fn=lambda x: x.body, full_event=True),
+            Complete(),
+            reducer,
+        ]
+    ).run()
 
     async def emit_valid_event(value):
         invocation_result = await controller.emit(value)
         assert invocation_result == value * 2
+
     time.sleep(flush_after_seconds + 0.2)  # Ensure different batch window
+
     async def emit_error_event(value):
         invocation_result = await controller.emit(value)
-        assert invocation_result == {'error': 'ValueError: Value -5 is negative!'}
+        assert invocation_result == {"error": "ValueError: Value -5 is negative!"}
 
     # Emit valid batch first (should succeed)
     tasks = [asyncio.create_task(emit_valid_event(v)) for v in [1, 2, 3]]  # All positive
@@ -2571,13 +2578,99 @@ async def emit_error_event(value):
     assert len(batch_result) == 6
 
     # First 3 should succeed
-    assert batch_result[0] == 2  # 1 * 2
-    assert batch_result[1] == 4  # 2 * 2
-    assert batch_result[2] == 6  # 3 * 2
+    assert batch_result[0] == 2
+    assert batch_result[1] == 4
+    assert batch_result[2] == 6
 
     # Next 3 should all have error (error propagated to all in batch)
     for single_result in batch_result[3:]:
-        assert single_result == {'error': 'ValueError: Value -5 is negative!'}
+        assert single_result == {"error": "ValueError: Value -5 is negative!"}
+
+
+def test_batch_error_handling_multiple_runnables():
+    asyncio.run(async_test_batch_error_handling_multiple_runnables())
+
+
+@pytest.mark.asyncio
+async def async_test_batch_error_handling_multiple_runnables():
+    """Test error handling in batched parallel execution with multiple runnables."""
+    flush_after_seconds = 0.3
+
+    class RunnableAddTen(ParallelExecutionRunnable):
+        def run(self, data, path, origin_name=None):
+            if isinstance(data, list):
+                return [item + 10 for item in data]
+            return data + 10
+
+    batch_size = 3
+    runnables = [
+        RunnableRaiseIfNegative("check_positive", raise_exception=False),
+        RunnableAddTen("add_ten", raise_exception=False),
+    ]
+
+    source = AsyncEmitSource()
+    batch_step = Batch(max_events=batch_size, full_event=True, flush_after_seconds=flush_after_seconds)
+    parallel_execution = ParallelExecution(
+        runnables,
+        execution_mechanism_by_runnable_name={
+            "check_positive": "naive",
+            "add_ten": "naive",
+        },
+    )
+    reducer = Reduce([], append_and_return)
+
+    controller = build_flow(
+        [
+            source,
+            batch_step,
+            parallel_execution,
+            FlatMap(fn=lambda x: x.body, full_event=True),
+            Complete(),
+            reducer,
+        ]
+    ).run()
+
+    async def emit_valid_event(value):
+        invocation_result = await controller.emit(value)
+        assert invocation_result["check_positive"] == value * 2
+        assert invocation_result["add_ten"] == value + 10
+
+    time.sleep(flush_after_seconds + 0.2)  # Ensure different batch window
+
+    async def emit_error_event(value):
+        invocation_result = await controller.emit(value)
+        # check_positive should have error, but add_ten should still work
+        assert invocation_result["check_positive"] == {"error": "ValueError: Value -5 is negative!"}
+        assert invocation_result["add_ten"] == value + 10
+
+    # Emit valid batch first (should succeed)
+    tasks = [asyncio.create_task(emit_valid_event(v)) for v in [1, 2, 3]]  # All positive
+    await asyncio.gather(*tasks)
+
+    # Emit batch with negative value concurrently (should error for check_positive, but add_ten works)
+    tasks = [asyncio.create_task(emit_error_event(v)) for v in [4, -5, 6]]  # Middle value is negative
+    await asyncio.gather(*tasks)
+
+    await controller.terminate()
+    batch_result = await controller.await_termination()
+
+    # Should have 6 results total (3 valid + 3 with partial error)
+    assert len(batch_result) == 6
+
+    # First 3 should succeed with both runnables
+    assert batch_result[0]["check_positive"] == 2  # 1 * 2
+    assert batch_result[0]["add_ten"] == 11  # 1 + 10
+    assert batch_result[1]["check_positive"] == 4  # 2 * 2
+    assert batch_result[1]["add_ten"] == 12  # 2 + 10
+    assert batch_result[2]["check_positive"] == 6  # 3 * 2
+    assert batch_result[2]["add_ten"] == 13  # 3 + 10
+
+    # Next 3 should have error for check_positive but add_ten should work
+    for i, single_result in enumerate(batch_result[3:]):
+        assert single_result["check_positive"] == {"error": "ValueError: Value -5 is negative!"}
+        # add_ten should still work for values [4, -5, 6]
+        expected_add_values = [14, 5, 16]
+        assert single_result["add_ten"] == expected_add_values[i]
 
 
 async def async_test_write_csv(tmpdir):