diff --git a/spec/namespaces/ml.yaml b/spec/namespaces/ml.yaml
index f7b860128..eaf1d07ae 100644
--- a/spec/namespaces/ml.yaml
+++ b/spec/namespaces/ml.yaml
@@ -147,7 +147,7 @@ components:
       content:
         application/json:
           schema:
-            $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+            $ref: '../schemas/_common.yaml#/components/schemas/WriteResponseBase'
     ml.register_model@200:
       content:
         application/json:
@@ -157,7 +157,7 @@ components:
       content:
         application/json:
           schema:
-            $ref: '../schemas/ml._common.yaml#/components/schemas/ModelGroup'
+            $ref: '../schemas/_common.yaml#/components/schemas/WriteResponseBase'
     ml.get_task@200:
       content:
         application/json:
diff --git a/spec/schemas/ml._common.yaml b/spec/schemas/ml._common.yaml
index 3e06e2612..e5bfacd2b 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.
-        state:
+        task_id:
+          type: string
+        status:
           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:
@@ -139,4 +140,4 @@ components:
         is_async:
           type: boolean
       required:
-        - state
+        - status
diff --git a/tests/default/ml/models.yaml b/tests/default/ml/models.yaml
index ac386c0a4..b650f5ba9 100644
--- a/tests/default/ml/models.yaml
+++ b/tests/default/ml/models.yaml
@@ -11,6 +11,8 @@ prologues:
           plugins:
             ml_commons:
               only_run_on_ml_node: false
+  - path: /_plugins/_ml/tasks/*
+    method: DELETE
 chapters:
   - synopsis: Create model.
     id: create_model
@@ -26,6 +28,7 @@ chapters:
     output:
       task_id: payload.task_id
   - synopsis: Wait for task.
+    id: register_model
     path: /_plugins/_ml/tasks/{task_id}
     method: GET
     warnings:
@@ -36,6 +39,13 @@ chapters:
       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: ${register_model.model_id}
+    method: DELETE