diff --git a/darwin/future/meta/objects/stage.py b/darwin/future/meta/objects/stage.py index 44decb650..e6411fd88 100644 --- a/darwin/future/meta/objects/stage.py +++ b/darwin/future/meta/objects/stage.py @@ -4,7 +4,7 @@ from uuid import UUID from darwin.future.core.items import get_item_ids_stage, move_items_to_stage -from darwin.future.data_objects.workflow import WFStageCore +from darwin.future.data_objects.workflow import WFEdgeCore, WFStageCore from darwin.future.meta.objects.base import MetaBase @@ -64,9 +64,6 @@ def type(self) -> str: return self._element.type.value @property - def edges(self) -> List[List[UUID]]: + def edges(self) -> List[WFEdgeCore]: """Edge ID, source stage ID, target stage ID.""" - return [ - [edge.id, edge.source_stage_id, edge.target_stage_id] - for edge in self._element.edges - ] + return list(self._element.edges) diff --git a/darwin/future/tests/meta/objects/test_stagemeta.py b/darwin/future/tests/meta/objects/test_stagemeta.py index 696d42cda..71b5035d4 100644 --- a/darwin/future/tests/meta/objects/test_stagemeta.py +++ b/darwin/future/tests/meta/objects/test_stagemeta.py @@ -129,8 +129,14 @@ def test_get_stage_edges(stage_meta): }, ) assert len(test_stage.edges) == 2 - assert len(test_stage.edges[0]) == 3 - assert test_stage.edges[0][0] == UUID("00000000-0000-0000-0000-000000000000") + assert test_stage.edges[0].name == "edge_1" + assert test_stage.edges[0].id == UUID("00000000-0000-0000-0000-000000000000") + assert test_stage.edges[0].source_stage_id == UUID( + "00000000-0000-0000-0000-000000000000" + ) + assert test_stage.edges[0].target_stage_id == UUID( + "00000000-0000-0000-0000-000000000000" + ) def test_stage_str_method(stage_meta: Stage) -> None: