From 6f7acf84293278b5e5577b515123c480c20ce4b9 Mon Sep 17 00:00:00 2001
From: Thomas Farr <tsfarr@amazon.com>
Date: Thu, 27 Feb 2025 13:40:19 +1300
Subject: [PATCH] Fix k-NN model training tests to work with small training
 sample set

Signed-off-by: Thomas Farr <tsfarr@amazon.com>
---
 tests/default/knn/models.yaml        | 13 +++++++++++++
 tests/default/knn/models/search.yaml | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/tests/default/knn/models.yaml b/tests/default/knn/models.yaml
index 3589fc162..4cfa893a5 100644
--- a/tests/default/knn/models.yaml
+++ b/tests/default/knn/models.yaml
@@ -27,6 +27,10 @@ prologues:
         - {recommendation_vector: [1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], duration: 12.2}
         - {index: {_index: movies, _id: '2'}}
         - {recommendation_vector: [2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5], duration: 7.1}
+        - {index: {_index: movies, _id: '3'}}
+        - {recommendation_vector: [3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5], duration: 4.8}
+        - {index: {_index: movies, _id: '4'}}
+        - {recommendation_vector: [4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5], duration: 18.6}
     status: [200]
   - method: POST
     path: /_plugins/_knn/models/model-1/_train
@@ -41,6 +45,15 @@ prologues:
         mode: on_disk
         compression_level: 32x
         spaceType: l2
+        method:
+          name: ivf
+          engine: faiss
+          parameters:
+            nlist: 4
+            encoder:
+              name: pq
+              parameters:
+                code_size: 2
     status: [200]
 epilogues:
   - path: /movies
diff --git a/tests/default/knn/models/search.yaml b/tests/default/knn/models/search.yaml
index 52711b0f5..d608b878e 100644
--- a/tests/default/knn/models/search.yaml
+++ b/tests/default/knn/models/search.yaml
@@ -27,6 +27,10 @@ prologues:
         - {recommendation_vector: [1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], duration: 12.2}
         - {index: {_index: movies, _id: '2'}}
         - {recommendation_vector: [2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5, 2.5], duration: 7.1}
+        - {index: {_index: movies, _id: '3'}}
+        - {recommendation_vector: [3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5, 3.5], duration: 4.8}
+        - {index: {_index: movies, _id: '4'}}
+        - {recommendation_vector: [4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5], duration: 18.6}
     status: [200]
   - method: POST
     path: /_plugins/_knn/models/model-1/_train
@@ -41,6 +45,15 @@ prologues:
         mode: on_disk
         compression_level: 32x
         spaceType: l2
+        method:
+          name: ivf
+          engine: faiss
+          parameters:
+            nlist: 4
+            encoder:
+              name: pq
+              parameters:
+                code_size: 2
     status: [200]
 epilogues:
   - path: /movies