Workflow Version missing overload (#2044)

cognitedata · Nov 23, 2024 · a7aa775 · a7aa775
1 parent 34a3338
commit a7aa775
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/cognite/client/_api/workflows.py b/cognite/client/_api/workflows.py
@@ -697,6 +697,24 @@ def delete(
             wrap_ids=True,
         )
 
+    @overload
+    def retrieve(
+        self,
+        workflow_external_id: WorkflowVersionIdentifier | str,
+        version: str | None = None,
+        *,
+        ignore_unknown_ids: bool = False,
+    ) -> WorkflowVersion | None: ...
+
+    @overload
+    def retrieve(
+        self,
+        workflow_external_id: Sequence[WorkflowVersionIdentifier] | WorkflowIds,
+        version: None = None,
+        *,
+        ignore_unknown_ids: bool = False,
+    ) -> WorkflowVersionList: ...
+
     def retrieve(
         self,
         workflow_external_id: WorkflowVersionIdentifier | Sequence[WorkflowVersionIdentifier] | WorkflowIds | str,

diff --git a/tests/tests_integration/test_api/test_workflows.py b/tests/tests_integration/test_api/test_workflows.py
@@ -373,10 +373,10 @@ def test_delete_non_existing(self, cognite_client: CogniteClient, new_workflow_v
             ],
             ignore_unknown_ids=True,
         )
-        assert cognite_client.workflows.versions.retrieve(*new_workflow_version.as_id().as_primitive()) is None
+        assert cognite_client.workflows.versions.retrieve(new_workflow_version.as_id().as_tuple()) is None
 
     def test_retrieve_workflow(self, cognite_client: CogniteClient, new_workflow_version: WorkflowVersion) -> None:
-        retrieved = cognite_client.workflows.versions.retrieve(*new_workflow_version.as_id().as_primitive())
+        retrieved = cognite_client.workflows.versions.retrieve(new_workflow_version.as_id().as_tuple())
         assert retrieved == new_workflow_version
 
     def test_retrieve_non_existing_workflow(self, cognite_client: CogniteClient) -> None: