From 80625b8702c17b54a5e0335c9d0d88d34ac94cf3 Mon Sep 17 00:00:00 2001
From: Nathalie Jonathan <nathhjo@amazon.com>
Date: Wed, 13 Nov 2024 17:05:24 -0800
Subject: [PATCH] Add missing tests for ML APIs (search and delete)

---
 spec/schemas/ml._common.yaml    | 91 ++++++++++++++++++++++++++++++++-
 tests/plugins/ml/ml/models.yaml | 25 +++++++++
 2 files changed, 114 insertions(+), 2 deletions(-)

diff --git a/spec/schemas/ml._common.yaml b/spec/schemas/ml._common.yaml
index 46cffc9f3..a36efaa28 100644
--- a/spec/schemas/ml._common.yaml
+++ b/spec/schemas/ml._common.yaml
@@ -22,6 +22,28 @@ components:
     SearchModelsResponse:
       type: object
       properties:
+        took:
+          type: integer
+          description: The time taken to execute the search.
+        timed_out:
+          type: boolean
+          description: Whether the search timed out.
+        _shards:
+          type: object
+          description: The shards information.
+          properties:
+            total:
+              type: integer
+              description: The total number of shards.
+            successful:
+              type: integer
+              description: The number of successful shards.
+            skipped:
+              type: integer
+              description: The number of skipped shards.
+            failed:
+              type: integer
+              description: The number of failed shards.
         hits:
           $ref: '#/components/schemas/SearchModelsHits'
       required:
@@ -35,6 +57,9 @@ components:
           type: array
           items:
             $ref: '#/components/schemas/SearchModelsHitsHit'
+        max_score:
+          type: number
+          description: The maximum score.
       required:
         - hits
         - total
@@ -59,12 +84,74 @@ components:
         _id:
           type: string
           description: The hit ID.
+        _version:
+          type: integer
+          description: The version.
+        _seq_no:
+          type: integer
+          description: The sequence number.
+        _primary_term:
+          type: integer
+          description: The primary term.
+        _score:
+          type: number
+          description: The score.
+        _source:
+          $ref: '#/components/schemas/Source'
         model_id:
           type: string
           description: The model ID.
       required:
-        - _id
-        - model_id
+        - _score
+    Source:
+      type: object
+      properties:
+        last_deployed_time:
+          type: number
+          description: The last deployed time.
+        model_version:
+          type: string
+          description: The model version.
+        created_time:
+          type: number
+          description: The created time.
+        deploy_to_all_nodes:
+          type: boolean
+          description: Whether to deploy to all nodes.
+        is_hidden:
+          type: boolean
+          description: Whether the model is hidden.
+        description:
+          type: string
+          description: The model description.
+        model_state:
+          type: string
+          description: The model state.
+        planning_worker_node_count:
+          type: number
+          description: The planning worker node count.
+        last_updated_time:
+          type: number
+          description: The last updated time.
+        name:
+          type: string
+          description: The model name.
+        connector_id:
+          type: string
+          description: The connector ID.
+        current_worker_node_count:
+          type: number
+          description: The current worker node count.
+        model_group_id:
+          type: string
+          description: The model group ID.
+        planning_worker_nodes:
+          type: array
+          items:
+            $ref: '_common.yaml#/components/schemas/NodeIds'
+        algorithm:
+          type: string
+          description: The algorithm.
     ModelGroupRegistration:
       type: object
       properties:
diff --git a/tests/plugins/ml/ml/models.yaml b/tests/plugins/ml/ml/models.yaml
index ce087654f..a52b0b3c9 100644
--- a/tests/plugins/ml/ml/models.yaml
+++ b/tests/plugins/ml/ml/models.yaml
@@ -22,6 +22,23 @@ chapters:
       status: 200
     output:
       task_id: payload.task_id
+  - synopsis: Search model.
+    id: search_model
+    path: /_plugins/_ml/models/_search
+    method: GET
+    request:
+      payload:
+        query:
+          match_all: {}
+        size: 1000
+    response:
+      status: 200
+      payload:
+        hits:
+          hits:
+            - _score: 1
+    output:
+      search_id: payload.hits.hits[0]._id
   - synopsis: Wait to get completed task.
     id: get_completed_task
     path: /_plugins/_ml/tasks/{task_id}
@@ -37,6 +54,14 @@ chapters:
     retry:
       count: 5
       wait: 30000
+  - synopsis: Delete task.
+    id: delete_task
+    path: /_plugins/_ml/tasks/{task_id}
+    method: DELETE
+    parameters:
+      task_id: ${register_model.task_id}
+    response:
+      status: 200
   - synopsis: Delete model.
     path: /_plugins/_ml/models/{model_id}
     parameters: