Support using multiple/concurrent EventListeners

griptape-ai · Sep 11, 2024 · 9781644 · 9781644
1 parent 5f0ff21
commit 9781644
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/griptape/events/event_bus.py b/griptape/events/event_bus.py
@@ -24,8 +24,9 @@ def add_event_listeners(self, event_listeners: list[EventListener]) -> list[Even
         return [self.add_event_listener(event_listener) for event_listener in event_listeners]
 
     def remove_event_listeners(self, event_listeners: list[EventListener]) -> None:
-        for event_listener in event_listeners:
-            self.remove_event_listener(event_listener)
+        with self._thread_lock:
+            for event_listener in event_listeners:
+                self.remove_event_listener(event_listener)
 
     def add_event_listener(self, event_listener: EventListener) -> EventListener:
         with self._thread_lock:

diff --git a/tests/unit/events/test_event_listener.py b/tests/unit/events/test_event_listener.py
@@ -94,8 +94,8 @@ def test_add_remove_event_listener(self, pipeline):
         mock1 = Mock()
         mock2 = Mock()
         # duplicate event listeners will only get added once
-        event_listener_1 = EventBus.add_event_listener(EventListener(mock1, event_types=[StartPromptEvent]))
-        EventBus.add_event_listener(EventListener(mock1, event_types=[StartPromptEvent]))
+        event_listener_1 = EventBus.add_event_listener(EventListener(mock1, id="1", event_types=[StartPromptEvent]))
+        EventBus.add_event_listener(EventListener(mock1, id="1", event_types=[StartPromptEvent]))
 
         event_listener_3 = EventBus.add_event_listener(EventListener(mock1, event_types=[FinishPromptEvent]))
         event_listener_4 = EventBus.add_event_listener(EventListener(mock2, event_types=[StartPromptEvent]))