diff --git a/spec/namespaces/ml.yaml b/spec/namespaces/ml.yaml
index eaf1d07ae..1cf5d989b 100644
--- a/spec/namespaces/ml.yaml
+++ b/spec/namespaces/ml.yaml
@@ -152,7 +152,15 @@ components:
       content:
         application/json:
           schema:
-            $ref: '../schemas/ml._common.yaml#/components/schemas/Task'
+            type: object
+            properties:
+              task_id:
+                type: string
+              status:
+                type: string
+            required:
+              - task_id
+              - status
     ml.delete_model@200:
       content:
         application/json:
diff --git a/tests/default/ml/models.yaml b/tests/default/ml/models.yaml
index 4e4a5e743..4bdac64e1 100644
--- a/tests/default/ml/models.yaml
+++ b/tests/default/ml/models.yaml
@@ -12,8 +12,8 @@ prologues:
             ml_commons:
               only_run_on_ml_node: false
 chapters:
-  - synopsis: Create model.
-    id: create_model
+  - synopsis: Register model.
+    id: register_model
     path: /_plugins/_ml/models/_register
     method: POST
     request:
@@ -25,14 +25,14 @@ chapters:
       status: 200
     output:
       task_id: payload.task_id
-  - synopsis: Wait for task.
-    id: register_model
+  - synopsis: Wait to get completed task.
+    id: get_completed_task
     path: /_plugins/_ml/tasks/{task_id}
     method: GET
     warnings:
       multiple-paths-detected: false
     parameters:
-      task_id: ${create_model.task_id}
+      task_id: ${register_model.task_id}
     response:
       status: 200
       payload:
@@ -45,5 +45,5 @@ chapters:
   - synopsis: Delete model.
     path: /_plugins/_ml/models/{model_id}
     parameters:
-      model_id: ${register_model.model_id}
+      model_id: ${get_completed_task.model_id}
     method: DELETE