From fce9d576e1c960aa38e598ac4b85ed376f3d4fc6 Mon Sep 17 00:00:00 2001
From: Neetika Singhal <neetiks@amazon.com>
Date: Tue, 19 Mar 2024 16:22:19 -0700
Subject: [PATCH] Fix SimpleNestedExplainIT.testExplainMultipleDocs flakiness

Signed-off-by: Neetika Singhal <neetiks@amazon.com>
---
 .../search/nested/SimpleNestedExplainIT.java  |  7 ++-
 .../test/OpenSearchIntegTestCase.java         | 43 ++++++++++++++-----
 2 files changed, 38 insertions(+), 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
index a6554271a0bc5..00835db05bfa4 100644
--- 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
index 7cb1b3f4fe0d8..281385a4fd034 100644
--- 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();