diff --git a/references/sdk.md b/references/sdk.md
index 6814d2c6..ceb95b0e 100644
--- a/references/sdk.md
+++ b/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
diff --git a/references/sdk_models.md b/references/sdk_models.md
index 40fecad7..0a9b70a9 100644
--- a/references/sdk_models.md
+++ b/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'>
 ```
diff --git a/substra/sdk/client.py b/substra/sdk/client.py
index 82bc273b..074b1686 100644
--- a/substra/sdk/client.py
+++ b/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,