diff --git a/storey/flow.py b/storey/flow.py index 5c7278c3..1ebf9bf2 100644 --- a/storey/flow.py +++ b/storey/flow.py @@ -111,7 +111,7 @@ def is_batched_event(event) -> bool: not isinstance(event, StreamCompletion) and isinstance(getattr(event, "body", None), list) and event.body - and any(hasattr(sub_event, "body") for sub_event in event.body) + and any(hasattr(sub_event, "body") and "event" in type(sub_event).__name__.lower() for sub_event in event.body) ) diff --git a/tests/test_flow.py b/tests/test_flow.py index c6d11d93..ed1013dc 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -81,6 +81,7 @@ RunnableExecutor, _Batching, _ConcurrentJobExecution, + is_batched_event, ) from tests.helpers import MockContext, MockLogger @@ -102,6 +103,20 @@ def raise_ex(self, element): return element +def test_is_batched_event_with_event_sub_events(): + batched = Event(body=[Event(body=1), Event(body=2)]) + assert is_batched_event(batched) is True + + +def test_is_batched_event_excludes_non_event_body_objects(): + class Response: + def __init__(self, body): + self.body = body + + batched = Event(body=[Response(body=1), Response(body=2)]) + assert is_batched_event(batched) is False + + def test_functional_flow(): controller = build_flow( [