diff --git a/spec/namespaces/ml.yaml b/spec/namespaces/ml.yaml
index 54b5ea532..2822ceb89 100644
--- a/spec/namespaces/ml.yaml
+++ b/spec/namespaces/ml.yaml
@@ -130,36 +130,52 @@ components:
     ml.search_models:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/SearchModelsQuery'
+          schema:
+            $ref: '../schemas/ml._common.yaml#/components/schemas/SearchModelsQuery'
   responses:
     ml.register_model_group@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroupRegistration'
+          schema:
+            $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroupRegistration'
     ml.get_model_group@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+          schema:
+            $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
     ml.delete_model_group@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+          schema:
+            $ref: '../schemas/_common.yaml#/components/schemas/WriteResponseBase'
     ml.register_model@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/Task'
+          schema:
+            type: object
+            properties:
+              status:
+                type: string
+              task_id:
+                type: string
+            required:
+              - status
+              - task_id
     ml.delete_model@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+          schema:
+            $ref: '../schemas/_common.yaml#/components/schemas/WriteResponseBase'
     ml.get_task@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/Task'
+          schema:
+            $ref: '../schemas/ml._common.yaml#/components/schemas/Task'
     ml.search_models@200:
       content:
         application/json:
-          $ref: '../schemas/ml._common.yaml#/components/schemas/SearchModelsResponse'
+          schema:
+            $ref: '../schemas/ml._common.yaml#/components/schemas/SearchModelsResponse'
   parameters:
     ml.get_model_group::path.model_group_id:
       name: model_group_id
diff --git a/spec/schemas/ml._common.yaml b/spec/schemas/ml._common.yaml
index 3e06e2612..8d1b40d3e 100644
--- a/spec/schemas/ml._common.yaml
+++ b/spec/schemas/ml._common.yaml
@@ -108,10 +108,10 @@ components:
       properties:
         model_id:
           type: string
-          description: The model ID.
+        task_id:
+          type: string
         state:
           type: string
-          description: The state.
           enum:
             - CANCELLED
             - COMPLETED
@@ -124,6 +124,7 @@ components:
           description: Task type.
           enum:
             - DEPLOY_MODEL
+            - REGISTER_MODEL
         function_name:
           type: string
         worker_node:
diff --git a/tests/default/ml/models.yaml b/tests/default/ml/models.yaml
index ac386c0a4..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,17 +25,25 @@ chapters:
       status: 200
     output:
       task_id: payload.task_id
-  - synopsis: Wait for task.
+  - 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:
         state: COMPLETED
+    output:
+      model_id: payload.model_id
     retry:
-      count: 3
+      count: 5
       wait: 30000
+  - synopsis: Delete model.
+    path: /_plugins/_ml/models/{model_id}
+    parameters:
+      model_id: ${get_completed_task.model_id}
+    method: DELETE