update workflow meta object - raising MissingDataset when dataset s…

…tage not found

darwin/future/meta/objects/workflow.py

@@ -7,6 +7,7 @@
 from darwin.dataset.upload_manager import LocalFile
 from darwin.datatypes import PathLike
 from darwin.future.data_objects.workflow import WFDatasetCore, WFTypeCore, WorkflowCore
+from darwin.future.exceptions import MissingDataset
 from darwin.future.meta.objects.base import MetaBase
 from darwin.future.meta.queries.stage import Stage, StageQuery
 
@@ -55,12 +56,14 @@ def stages(self) -> StageQuery:
             meta_params["dataset_name"] = self.datasets[0].name
         return StageQuery(self.client, meta_params=meta_params)
 
-    def _get_dataset_stage(self) -> Stage | None:
+    def _get_dataset_stage(self) -> Stage:
         # stages are not in right order - finding the dataset stage
         for stage in self.stages:
             if stage.type == "dataset":
                 return stage
 
+        raise MissingDataset(f"Workflow has no dataset stage")
+
     @property
     def datasets(self) -> List[WFDatasetCore]:
         if self._element.dataset is None:
@@ -81,13 +84,11 @@ def push_from_dataset_stage(self, wait: bool = True) -> Workflow:
         assert len(stages) > 1
 
         ds_stage = self._get_dataset_stage()
-        assert ds_stage
-
         assert ds_stage._element.type == WFTypeCore.DATASET
         next_stage = ds_stage._element.edges[0].target_stage_id
         assert next_stage is not None
-
         ds_stage.move_attached_files_to_stage(next_stage, wait)
+
         return self
 
     def upload_files(