Fix tests

griptape-ai · Jun 4, 2024 · c8bc783 · c8bc783
1 parent b18108f
commit c8bc783
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/tests/unit/drivers/event_listener/test_pusher_event_listener_driver.py b/tests/unit/drivers/event_listener/test_pusher_event_listener_driver.py
@@ -28,15 +28,15 @@ def test_init(self, driver):
         assert driver
 
     def test_try_publish_event_payload(self, driver):
-        driver.try_publish_event_payload(MockEvent().to_dict())
+        data = MockEvent().to_dict()
+        driver.try_publish_event_payload(data)
 
-        assert driver.pusher_client.trigger.called_with(
-            channels="test-channel", event_name="test-event", data=MockEvent().to_dict()
-        )
+        driver.pusher_client.trigger.assert_called_with(channels="test-channel", event_name="test-event", data=data)
 
     def test_try_publish_event_payload_batch(self, driver):
-        driver.try_publish_event_payload_batch([MockEvent().to_dict() for _ in range(3)])
+        data = [MockEvent().to_dict() for _ in range(3)]
+        driver.try_publish_event_payload_batch(data)
 
-        assert driver.pusher_client.trigger_batch.called_with(
-            [{"channel": "test-channel", "name": "test-event", "data": MockEvent().to_dict()} for _ in range(3)]
+        driver.pusher_client.trigger_batch.assert_called_with(
+            [{"channel": "test-channel", "name": "test-event", "data": data[i]} for i in range(3)]
         )