fix: wait_function (#400)

Signed-off-by: SdgJlbl <[email protected]>

references/sdk.md

@@ -644,6 +644,27 @@ It is considered finished when the status is done, failed or cancelled.
 
  - `exceptions.FutureTimeoutError`: The compute plan took more than the duration set in the timeout to complete.
 Not raised when `timeout == None`
+## wait_function
+```text
+wait_function(self, key: str, *, timeout: Optional[float] = None, polling_period: float = 1.0, raise_on_failure: bool = True) -> substra.sdk.models.Function
+```
+
+Wait for the build of the given function to finish.
+It is considered finished when the status is ready, failed or cancelled.
+
+**Arguments:**
+ - `key (str, required)`: the key of the task to wait for.
+ - `timeout (float, optional)`: maximum time to wait, in seconds. If set to None, will hang until completion.
+ - `polling_period (float, required)`: time to wait between two checks, in seconds. Defaults to 2.0.
+ - `raise_on_failure (bool, required)`: whether to raise an exception if the execution fails. Defaults to True.
+
+**Returns:**
+
+ - `models.Function`: the function once built
+Raises:
+   exceptions.FutureFailureError: The function build failed or have been cancelled.
+   exceptions.FutureTimeoutError: The function took more than the duration set in the timeout to build.
+       Not raised when `timeout == None`
 ## wait_task
 ```text
 wait_task(self, key: str, *, timeout: Optional[float] = None, polling_period: float = 2.0, raise_on_failure: bool = True) -> substra.sdk.models.Task

references/sdk_models.md

@@ -71,6 +71,7 @@ Asset creation specification base class.
 - creation_date: <class 'datetime.datetime'>
 - inputs: typing.List[substra.sdk.models.FunctionInput]
 - outputs: typing.List[substra.sdk.models.FunctionOutput]
+- status: <enum 'FunctionStatus'>
 - description: <class 'substra.sdk.models._File'>
 - archive: <class 'substra.sdk.models._File'>
 ```

substra/sdk/client.py

@@ -1079,7 +1079,7 @@ def wait_task(
     @logit
     def wait_function(
         self, key: str, *, timeout: Optional[float] = None, polling_period: float = 1.0, raise_on_failure: bool = True
-    ) -> models.Task:
+    ) -> models.Function:
         """Wait for the build of the given function to finish.
 
         It is considered finished when the status is ready, failed or cancelled.
@@ -1091,17 +1091,21 @@ def wait_function(
             raise_on_failure (bool): whether to raise an exception if the execution fails. Defaults to True.
 
         Returns:
-            models.Task: the task after completion
+            models.Function: the function once built
 
          Raises:
-            exceptions.FutureFailureError: The task failed or have been cancelled.
-            exceptions.FutureTimeoutError: The task took more than the duration set in the timeout to complete.
+            exceptions.FutureFailureError: The function build failed or have been cancelled.
+            exceptions.FutureTimeoutError: The function took more than the duration set in the timeout to build.
                 Not raised when `timeout == None`
         """
         asset_getter = self.get_function
         status_canceled = models.FunctionStatus.canceled.value
         status_failed = models.FunctionStatus.failed.value
-        statuses_stopped = (models.FunctionStatus.ready.value, models.FunctionStatus.canceled.value)
+        statuses_stopped = (
+            models.FunctionStatus.ready.value,
+            models.FunctionStatus.canceled.value,
+            models.FunctionStatus.failed.value,
+        )
         return self._wait(
             key=key,
             asset_getter=asset_getter,