update the indexRandom function to create more segments

Signed-off-by: Neetika Singhal <neetiks@amazon.com>
opensearch-project · Oct 2, 2023 · 9d153ad · 9d153ad
1 parent 6003560
commit 9d153ad
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java b/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java
@@ -1642,6 +1642,10 @@ public void indexRandom(boolean forceRefresh, boolean dummyDocuments, boolean ma
             }
         }
         assertThat(actualErrors, emptyIterable());
+
+        // keeping the default segment count to be 2
+        bogusIds.addAll(indexRandomForMultipleSlices(2, indicesArray));
+
         if (!bogusIds.isEmpty()) {
             // delete the bogus types again - it might trigger merges or at least holes in the segments and enforces deleted docs!
             for (List<String> doc : bogusIds) {
@@ -1659,6 +1663,28 @@ public void indexRandom(boolean forceRefresh, boolean dummyDocuments, boolean ma
         }
     }
 
+    /*
+    * This method ingests bogus documents for the given indices such that multiple slices
+    * are formed. This is greatly useful for testing with the concurrent search use-case.
+    * @param refreshCount    the number of segments to be formed by ingesting bogus documents
+    * @param indices         the indices in which bogus documents should be ingested
+    * */
+    public Set<List<String>> indexRandomForMultipleSlices(int segments, String... indices) {
+        Set<List<String>> bogusIds = new HashSet<>();
+        for (String index : indices) {
+            int numDocs = getNumShards(index).totalNumShards * segments;
+            while (segments-- > 0) {
+                for (int i = 0; i < numDocs; i++) {
+                    String id = "bogus_doc_" + randomRealisticUnicodeOfLength(between(1, 10)) + dummmyDocIdGenerator.incrementAndGet();
+                    client().prepareIndex().setIndex(index).setId(id).setSource("{}", MediaTypeRegistry.JSON).setRouting(id).get();
+                    bogusIds.add(Arrays.asList(index, id));
+                }
+                refresh();
+            }
+        }
+        return bogusIds;
+    }
+
     private final AtomicInteger dummmyDocIdGenerator = new AtomicInteger();
 
     /** Disables an index block for the specified index */

diff --git a/test/framework/src/main/java/org/opensearch/test/ParameterizedOpenSearchIntegTestCase.java b/test/framework/src/main/java/org/opensearch/test/ParameterizedOpenSearchIntegTestCase.java
@@ -13,6 +13,8 @@
 import org.junit.After;
 import org.junit.Before;
 
+import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
+
 /**
  * Base class for running the tests with parameterization of the dynamic settings
  * For any class that wants to use parameterization, use @ParametersFactory to generate
@@ -44,4 +46,14 @@ public void afterTests() {
         dynamicSettings.keySet().forEach(settingsToUnset::putNull);
         client().admin().cluster().prepareUpdateSettings().setPersistentSettings(settingsToUnset).get();
     }
+
+    public void indexRandomForConcurrentSearch(int segments, String... indices) {
+        if (dynamicSettings.get(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey()).equals("true")) {
+            indexRandomForMultipleSlices(segments, indices);
+        }
+    }
+
+    public void indexRandomForConcurrentSearch(String... indices) {
+        indexRandomForConcurrentSearch(2, indices);
+    }
 }