Fixes broken unit tests

DAGWorks-Inc · Aug 26, 2024 · eb6b21e · eb6b21e
1 parent 76fe23e
commit eb6b21e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 3 deletions.
diff --git a/burr/core/application.py b/burr/core/application.py
@@ -1358,7 +1358,7 @@ def callback(
                     "post_run_step",
                     app_id=self._uid,
                     partition_key=self._partition_key,
-                    action=next_action.name,
+                    action=next_action,
                     state=self._state,
                     result=result,
                     sequence_id=self.sequence_id,
@@ -1380,7 +1380,7 @@ def callback(
             if next_action.single_step:
                 next_action = cast(SingleStepStreamingAction, next_action)
                 generator = _run_single_step_streaming_action(
-                    action=next_action.name,
+                    action=next_action,
                     state=self._state,
                     inputs=action_inputs,
                     sequence_id=self.sequence_id,
@@ -1394,7 +1394,7 @@ def callback(
             else:
                 next_action = cast(StreamingAction, next_action)
                 generator = _run_multi_step_streaming_action(
-                    action=next_action.name,
+                    action=next_action,
                     state=self._state,
                     inputs=action_inputs,
                     sequence_id=self.sequence_id,

diff --git a/tests/core/test_application.py b/tests/core/test_application.py
@@ -2865,6 +2865,46 @@ def test_application_builder_initialize_fork_app_id_happy_pth():
 
 
 class NoOpTracker(SyncTrackingClient):
+    def copy(self):
+        pass
+
+    def pre_start_stream(
+        self,
+        *,
+        action: str,
+        sequence_id: int,
+        app_id: str,
+        partition_key: Optional[str],
+        **future_kwargs: Any,
+    ):
+        pass
+
+    def post_stream_item(
+        self,
+        *,
+        item: Any,
+        item_index: int,
+        stream_initialize_time: datetime.datetime,
+        first_stream_item_start_time: datetime.datetime,
+        action: str,
+        sequence_id: int,
+        app_id: str,
+        partition_key: Optional[str],
+        **future_kwargs: Any,
+    ):
+        pass
+
+    def post_end_stream(
+        self,
+        *,
+        action: str,
+        sequence_id: int,
+        app_id: str,
+        partition_key: Optional[str],
+        **future_kwargs: Any,
+    ):
+        pass
+
     def do_log_attributes(self, **future_kwargs: Any):
         pass