From acf77882dfc7eea77a12570a0f73da2d0419ff81 Mon Sep 17 00:00:00 2001
From: dblock <dblock@amazon.com>
Date: Tue, 5 Nov 2024 09:31:00 -0500
Subject: [PATCH] Fix: neural search error response.

Signed-off-by: dblock <dblock@amazon.com>
---
 spec/schemas/query._common.yaml               |  8 ++++++
 .../ml/ingest/pipeline/neural_search.yaml     | 17 ------------
 .../ml/ingest/pipeline/search_by_image.yaml   | 26 +++++++++++++++++++
 3 files changed, 34 insertions(+), 17 deletions(-)
 create mode 100644 tests/plugins/ml/ingest/pipeline/search_by_image.yaml

diff --git a/spec/schemas/query._common.yaml b/spec/schemas/query._common.yaml
index 40277b8b0..ef4e315f4 100644
--- a/spec/schemas/query._common.yaml
+++ b/spec/schemas/query._common.yaml
@@ -71,6 +71,14 @@ components:
         - error
 
     ErrorResponse:
+      type: object
+      properties:
+        error:
+          $ref: '#/components/schemas/Error'
+        status:
+          type: number
+
+    Error:
       type: object
       properties:
         root_cause:
diff --git a/tests/plugins/ml/ingest/pipeline/neural_search.yaml b/tests/plugins/ml/ingest/pipeline/neural_search.yaml
index ca2fdd60a..5e68f5995 100644
--- a/tests/plugins/ml/ingest/pipeline/neural_search.yaml
+++ b/tests/plugins/ml/ingest/pipeline/neural_search.yaml
@@ -183,23 +183,6 @@ chapters:
                       title: Moneyball
                   script:
                     source: _score * 1.7
-  - synopsis: Search by Image - Invalid Model.
-    path: /{index}/_search
-    method: POST
-    parameters:
-      index: movies
-    request:
-      payload:
-        _source:
-          excludes: [passage_embedding]
-        query:
-          neural:
-            passage_embedding:
-              query_image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=
-              model_id: invalid # This makes the test fail with a 404 but will still validate the query_image.
-              k: 100
-    response:
-      status: 404
   - synopsis: Undeploy a model.
     path: /_plugins/_ml/models/{model_id}/_undeploy
     method: POST
diff --git a/tests/plugins/ml/ingest/pipeline/search_by_image.yaml b/tests/plugins/ml/ingest/pipeline/search_by_image.yaml
new file mode 100644
index 000000000..e4fc46e61
--- /dev/null
+++ b/tests/plugins/ml/ingest/pipeline/search_by_image.yaml
@@ -0,0 +1,26 @@
+$schema: ../../../../../json_schemas/test_story.schema.yaml
+
+description: Test search by image.
+distributions:
+  excluded:
+    - amazon-managed
+    - amazon-serverless
+version: '>= 2.11'
+chapters:
+  - synopsis: Search by Image - Invalid Model.
+    path: /{index}/_search
+    method: POST
+    parameters:
+      index: movies
+    request:
+      payload:
+        _source:
+          excludes: [passage_embedding]
+        query:
+          neural:
+            passage_embedding:
+              query_image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=
+              model_id: invalid # This makes the test fail with a 404 but will still validate the query_image.
+              k: 100
+    response:
+      status: 404