From a364584c24bd0db0c596c5e862654fdef9e5a411 Mon Sep 17 00:00:00 2001
From: Suraj Singh <surajrider@gmail.com>
Date: Mon, 5 Jun 2023 16:51:26 -0700
Subject: [PATCH] Add basic integration test for segment replication

Signed-off-by: Suraj Singh <surajrider@gmail.com>
---
 .../opensearch/knn/index/OpenSearchIT.java    | 24 +++++++++++++++++++
 .../org/opensearch/knn/KNNRestTestCase.java   | 15 ++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/src/test/java/org/opensearch/knn/index/OpenSearchIT.java b/src/test/java/org/opensearch/knn/index/OpenSearchIT.java
index e4871e357b..bc9461d86f 100644
--- a/src/test/java/org/opensearch/knn/index/OpenSearchIT.java
+++ b/src/test/java/org/opensearch/knn/index/OpenSearchIT.java
@@ -55,6 +55,30 @@ public static void setUpClass() throws IOException {
         testData = new TestUtils.TestData(testIndexVectors.getPath(), testQueries.getPath());
     }
 
+    /**
+     * This test performs basic validation steps such as index, search in order to verify that kNN works with Segment
+     * replication.
+     * @throws Exception
+     */
+    public void testSegmentReplication() throws Exception {
+        String field1 = "field1";
+        String field2 = "field2";
+        createKnnIndex(INDEX_NAME, getKNNSegrepIndexSettings(), createKnnIndexMapping(Arrays.asList(field1, field2), Arrays.asList(2, 2)));
+
+        Float[] vector = { 6.0f, 7.0f };
+        addKnnDoc(INDEX_NAME, "1", Arrays.asList(field1, field2), Arrays.asList(vector, vector));
+
+        addKnnDoc(INDEX_NAME, "2", field1, vector);
+        addKnnDoc(INDEX_NAME, "3", field1, vector);
+        addKnnDoc(INDEX_NAME, "4", field1, vector);
+
+        addKnnDoc(INDEX_NAME, "5", field2, vector);
+        addKnnDoc(INDEX_NAME, "6", field2, vector);
+
+        refreshAllIndices();
+        assertEquals(6, getDocCount(INDEX_NAME));
+    }
+
     public void testEndToEnd() throws Exception {
         String indexName = "test-index-1";
         KNNEngine knnEngine1 = KNNEngine.NMSLIB;
diff --git a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
index a8203f2476..adf17fdd72 100644
--- a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
+++ b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
@@ -10,10 +10,12 @@
 import com.google.common.primitives.Floats;
 import org.apache.commons.lang.StringUtils;
 import org.apache.hc.core5.http.io.entity.EntityUtils;
+import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.Strings;
 import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.xcontent.XContentHelper;
 import org.opensearch.index.query.MatchAllQueryBuilder;
+import org.opensearch.indices.replication.common.ReplicationType;
 import org.opensearch.knn.index.query.KNNQueryBuilder;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.SpaceType;
@@ -97,6 +99,7 @@
 import static org.opensearch.knn.TestUtils.QUERY_VALUE;
 import static org.opensearch.knn.TestUtils.computeGroundTruthValues;
 
+import static org.opensearch.knn.index.KNNSettings.KNN_INDEX;
 import static org.opensearch.knn.index.SpaceType.L2;
 import static org.opensearch.knn.index.memory.NativeMemoryCacheManager.GRAPH_COUNT;
 import static org.opensearch.knn.index.util.KNNEngine.FAISS;
@@ -527,6 +530,18 @@ protected Settings getKNNDefaultIndexSettings() {
         return Settings.builder().put("number_of_shards", 1).put("number_of_replicas", 0).put("index.knn", true).build();
     }
 
+    /**
+     * Return default segment replication enabled index settings
+     */
+    protected Settings getKNNSegrepIndexSettings() {
+        return Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+            .put(IndexMetadata.SETTING_REPLICATION_TYPE, ReplicationType.SEGMENT)
+            .put(KNN_INDEX, true)
+            .build();
+    }
+
     /**
      * Get Stats from KNN Plugin
      */