Fix test

griptape-ai · Aug 8, 2024 · c391b14 · c391b14
1 parent 729f3aa
commit c391b14
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/unit/drivers/prompt/test_base_prompt_driver.py b/tests/unit/drivers/prompt/test_base_prompt_driver.py
@@ -1,7 +1,7 @@
 from griptape.artifacts import ErrorArtifact, TextArtifact
 from griptape.common import Message, PromptStack
 from griptape.events import FinishPromptEvent, StartPromptEvent
-from griptape.events.event_bus import _event_bus
+from griptape.events.event_bus import _EventBus
 from griptape.structures import Pipeline
 from griptape.tasks import PromptTask, ToolkitTask
 from tests.mocks.mock_failing_prompt_driver import MockFailingPromptDriver
@@ -27,7 +27,7 @@ def test_run_via_pipeline_retries_failure(self):
         assert isinstance(pipeline.run().output_task.output, ErrorArtifact)
 
     def test_run_via_pipeline_publishes_events(self, mocker):
-        mock_publish_event = mocker.patch.object(_event_bus, "publish_event")
+        mock_publish_event = mocker.patch.object(_EventBus, "publish_event")
         driver = MockPromptDriver()
         pipeline = Pipeline(prompt_driver=driver)
         pipeline.add_task(PromptTask("test"))