From acf7ee842d17ebdd829bda524846c6c4c7d2e168 Mon Sep 17 00:00:00 2001
From: dblock <dblock@amazon.com>
Date: Thu, 24 Oct 2024 15:54:39 -0400
Subject: [PATCH] Fix: model training intermittent test failure.

Signed-off-by: dblock <dblock@amazon.com>
---
 spec/namespaces/knn.yaml           |  6 +++++-
 tests/default/knn/train_model.yaml | 19 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/spec/namespaces/knn.yaml b/spec/namespaces/knn.yaml
index c38bba5b7..77b6a34cb 100644
--- a/spec/namespaces/knn.yaml
+++ b/spec/namespaces/knn.yaml
@@ -283,7 +283,11 @@ components:
       required: true
   responses:
     knn.delete_model@200: {}
-    knn.get_model@200: {}
+    knn.get_model@200:
+      content:
+        application/json:
+          schema:
+            type: object
     knn.search_models@200: {}
     knn.stats@200: {}
     knn.train_model@200:
diff --git a/tests/default/knn/train_model.yaml b/tests/default/knn/train_model.yaml
index b45ac5c6f..cb58ed97c 100644
--- a/tests/default/knn/train_model.yaml
+++ b/tests/default/knn/train_model.yaml
@@ -41,6 +41,9 @@ prologues:
         - {index: {_index: movies, _id: '9'}}
         - {recommendation_vector: [9.5, 9.5, 9.5, 9.5, 9.5, 9.5, 9.5, 9.5], duration: 8.9}
     status: [200]
+  - method: POST
+    path: /_refresh
+    status: [200]
 epilogues:
   - path: /movies
     method: DELETE
@@ -70,4 +73,18 @@ chapters:
     response:
       status: 200
     output:
-      test_model_id: payload.model_id
\ No newline at end of file
+      test_model_id: payload.model_id
+  - synopsis: Wait for the model to get trained.
+    warnings:
+      multiple-paths-detected: false
+    method: GET
+    path: /_plugins/_knn/models/{model_id}
+    parameters:
+      model_id: ${train_model.test_model_id}
+    retry:
+      count: 3
+    response:
+      status: 200
+      payload:
+        model_id: ${train_model.test_model_id}
+        state: created