Few changes in integration test (#1395) (#1404)

Signed-off-by: Heemin Kim <[email protected]> (cherry picked from commit 6abec19) Co-authored-by: Heemin Kim <[email protected]>
opensearch-project · Jan 22, 2024 · 4d9da8d · 4d9da8d
1 parent f41e44a
commit 4d9da8d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 20 deletions.
diff --git a/src/test/java/org/opensearch/knn/index/AdvancedFilteringUseCasesIT.java b/src/test/java/org/opensearch/knn/index/AdvancedFilteringUseCasesIT.java
@@ -448,14 +448,16 @@ public void testFiltering_whenNonNestedKNNAndNestedFilterAndNonNestedFieldWithNe
     private void validateFilterSearch(final String query, final String engine) throws IOException {
         String response = EntityUtils.toString(performSearch(INDEX_NAME, query).getEntity());
         // Validate number of documents returned as the expected number of documents
-        Assert.assertEquals("For engine " + engine + " : ", DOCUMENT_IN_RESPONSE, parseHits(response));
+        Assert.assertEquals("For engine " + engine + ", hits: ", DOCUMENT_IN_RESPONSE, parseHits(response));
+        Assert.assertEquals("For engine " + engine + ", totalSearchHits: ", k, parseTotalSearchHits(response));
         if (KNNEngine.getEngine(engine) == KNNEngine.FAISS) {
             // Update the filter threshold to 0 to ensure that we are hitting ANN Search use case for FAISS
             updateIndexSettings(INDEX_NAME, Settings.builder().put(KNNSettings.ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD, 0));
             response = EntityUtils.toString(performSearch(INDEX_NAME, query).getEntity());
 
             // Validate number of documents returned as the expected number of documents
-            Assert.assertEquals("For engine " + engine + " with ANN search :", DOCUMENT_IN_RESPONSE, parseHits(response));
+            Assert.assertEquals("For engine " + engine + ", hits with ANN search :", DOCUMENT_IN_RESPONSE, parseHits(response));
+            Assert.assertEquals("For engine " + engine + ", totalSearchHits with ANN search :", k, parseTotalSearchHits(response));
         }
     }
 

diff --git a/src/test/java/org/opensearch/knn/index/NestedSearchIT.java b/src/test/java/org/opensearch/knn/index/NestedSearchIT.java
@@ -63,17 +63,16 @@ public final void cleanUp() {
     public void testNestedSearchWithLucene_whenKIsTwo_thenReturnTwoResults() {
         createKnnIndex(2, KNNEngine.LUCENE.getName());
 
-        String doc1 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED)
-            .addVectors(FIELD_NAME_VECTOR, new Float[] { 1f, 1f }, new Float[] { 1f, 1f })
-            .build();
-        addKnnDoc(INDEX_NAME, "1", doc1);
-
-        String doc2 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED)
-            .addVectors(FIELD_NAME_VECTOR, new Float[] { 2f, 2f }, new Float[] { 2f, 2f })
-            .build();
-        addKnnDoc(INDEX_NAME, "2", doc2);
+        int totalDocCount = 15;
+        for (int i = 0; i < totalDocCount; i++) {
+            String doc = NestedKnnDocBuilder.create(FIELD_NAME_NESTED)
+                .addVectors(FIELD_NAME_VECTOR, new Float[] { (float) i, (float) i }, new Float[] { (float) i, (float) i })
+                .build();
+            addKnnDoc(INDEX_NAME, String.valueOf(i), doc);
+        }
 
         refreshIndex(INDEX_NAME);
+        forceMergeKnnIndex(INDEX_NAME);
 
         Float[] queryVector = { 1f, 1f };
         Response response = queryNestedField(INDEX_NAME, 2, queryVector);
@@ -86,17 +85,16 @@ public void testNestedSearchWithLucene_whenKIsTwo_thenReturnTwoResults() {
     public void testNestedSearchWithFaiss_whenKIsTwo_thenReturnTwoResults() {
         createKnnIndex(2, KNNEngine.FAISS.getName());
 
-        String doc1 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED)
-            .addVectors(FIELD_NAME_VECTOR, new Float[] { 1f, 1f }, new Float[] { 1f, 1f })
-            .build();
-        addKnnDoc(INDEX_NAME, "1", doc1);
-
-        String doc2 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED)
-            .addVectors(FIELD_NAME_VECTOR, new Float[] { 2f, 2f }, new Float[] { 2f, 2f })
-            .build();
-        addKnnDoc(INDEX_NAME, "2", doc2);
+        int totalDocCount = 15;
+        for (int i = 0; i < totalDocCount; i++) {
+            String doc = NestedKnnDocBuilder.create(FIELD_NAME_NESTED)
+                .addVectors(FIELD_NAME_VECTOR, new Float[] { (float) i, (float) i }, new Float[] { (float) i, (float) i })
+                .build();
+            addKnnDoc(INDEX_NAME, String.valueOf(i), doc);
+        }
 
         refreshIndex(INDEX_NAME);
+        forceMergeKnnIndex(INDEX_NAME);
 
         Float[] queryVector = { 1f, 1f };
         Response response = queryNestedField(INDEX_NAME, 2, queryVector);
@@ -148,6 +146,7 @@ public void testNestedSearchWithFaiss_whenDoingExactSearch_thenReturnCorrectResu
             addKnnDoc(INDEX_NAME, String.valueOf(i), doc);
         }
         refreshIndex(INDEX_NAME);
+        forceMergeKnnIndex(INDEX_NAME);
 
         // Make it as an exact search by setting the threshold larger than size of filteredIds(6)
         updateIndexSettings(INDEX_NAME, Settings.builder().put(KNNSettings.ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD, 100));