Fix SimpleNestedExplainIT.testExplainMultipleDocs flakiness

Signed-off-by: Neetika Singhal <[email protected]>
opensearch-project · Mar 19, 2024 · fce9d57 · fce9d57
1 parent d4e1ab1
commit fce9d57
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 12 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/search/nested/SimpleNestedExplainIT.java b/server/src/internalClusterTest/java/org/opensearch/search/nested/SimpleNestedExplainIT.java
@@ -14,6 +14,9 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.test.OpenSearchIntegTestCase;
 
+import java.util.List;
+import java.util.Set;
+
 import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.nestedQuery;
@@ -70,7 +73,7 @@ public void testExplainMultipleDocs() throws Exception {
             .setRefreshPolicy(IMMEDIATE)
             .get();
 
-        indexRandomForMultipleSlices("test");
+        Set<List<String>> bogusIds = indexRandomForMultipleSlices(false, "test");
 
         // Turn off the concurrent search setting to test search with non-concurrent search
         client().admin()
@@ -111,5 +114,7 @@ public void testExplainMultipleDocs() throws Exception {
             .prepareUpdateSettings()
             .setPersistentSettings(Settings.builder().putNull(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey()).build())
             .get();
+
+        deleteBogusDocs(bogusIds);
     }
 }
diff --git a/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java b/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java
@@ -1574,13 +1574,29 @@ public void indexRandom(boolean forceRefresh, boolean dummyDocuments, boolean ma
         }
     }
 
+    /*
+     * This method deletes the given bogus documents
+     * @param bogusIds doc ids to be deleted
+     */
+    protected void deleteBogusDocs(Set<List<String>> bogusIds) {
+        for (List<String> doc : bogusIds) {
+            assertEquals(
+                "failed to delete a dummy doc [" + doc.get(0) + "][" + doc.get(1) + "]",
+                DocWriteResponse.Result.DELETED,
+                client().prepareDelete(doc.get(0), doc.get(1)).setRouting(doc.get(1)).get().getResult()
+            );
+        }
+    }
+
     /*
      * This method ingests bogus documents for the given indices such that multiple slices
      * are formed. This is useful for testing with the concurrent search use-case as it creates
-     * multiple slices based on segment count.
-     * @param indices         the indices in which bogus documents should be ingested
+     * multiple slices based on segment count. When calling this method with value of toDeleteBogusDocs
+     * as true, don't forget to call deleteBogusDocs at the end of the test to delete the bogus docs
+     * @param toDeleteBogusDocs  provides an option to delete the bogus docs optionally
+     * @param indices            the indices in which bogus documents should be ingested
      * */
-    protected void indexRandomForMultipleSlices(String... indices) throws InterruptedException {
+    protected Set<List<String>> indexRandomForMultipleSlices(boolean toDeleteBogusDocs, String... indices) throws InterruptedException {
         Set<List<String>> bogusIds = new HashSet<>();
         int refreshCount = randomIntBetween(2, 3);
         for (String index : indices) {
@@ -1617,15 +1633,20 @@ protected void indexRandomForMultipleSlices(String... indices) throws Interrupte
                 refresh(index);
             }
         }
-        for (List<String> doc : bogusIds) {
-            assertEquals(
-                "failed to delete a dummy doc [" + doc.get(0) + "][" + doc.get(1) + "]",
-                DocWriteResponse.Result.DELETED,
-                client().prepareDelete(doc.get(0), doc.get(1)).setRouting(doc.get(1)).get().getResult()
-            );
+        if (toDeleteBogusDocs) {
+            deleteBogusDocs(bogusIds);
+            // refresh is called to make sure the bogus docs doesn't affect the search results
+            refresh();
         }
-        // refresh is called to make sure the bogus docs doesn't affect the search results
-        refresh();
+        return bogusIds;
+    }
+
+    /*
+    * Default behavior of indexRandomForMultipleSlices is to always delete the bogus docs in the
+    * method itself to avoid the need to be done by the caller explicitly
+    * */
+    protected void indexRandomForMultipleSlices(String... indices) throws InterruptedException {
+        indexRandomForMultipleSlices(true, indices);
     }
 
     private final AtomicInteger dummmyDocIdGenerator = new AtomicInteger();