From 215b9365030e83970d72b26fd2ba1528f2eb56e8 Mon Sep 17 00:00:00 2001
From: Shourya Dutta Biswas <114977491+shourya035@users.noreply.github.com>
Date: Fri, 1 Sep 2023 17:07:15 +0530
Subject: [PATCH] Adding assertBusy to SearchSliceIT

Signed-off-by: Shourya Dutta Biswas <114977491+shourya035@users.noreply.github.com>
---
 .../search/slice/SearchSliceIT.java           | 55 ++++++++++---------
 1 file changed, 29 insertions(+), 26 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/search/slice/SearchSliceIT.java b/server/src/internalClusterTest/java/org/opensearch/search/slice/SearchSliceIT.java
index d50f750a2b2ec..47e14b7c4b191 100644
--- a/server/src/internalClusterTest/java/org/opensearch/search/slice/SearchSliceIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/search/slice/SearchSliceIT.java
@@ -57,6 +57,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
@@ -299,37 +300,39 @@ public void testInvalidQuery() throws Exception {
         assertThat(rootCause.getMessage(), equalTo("`slice` cannot be used outside of a scroll context or PIT context"));
     }
 
-    private void assertSearchSlicesWithScroll(SearchRequestBuilder request, String field, int numSlice, int numDocs) {
-        int totalResults = 0;
-        List<String> keys = new ArrayList<>();
-        for (int id = 0; id < numSlice; id++) {
-            SliceBuilder sliceBuilder = new SliceBuilder(field, id, numSlice);
-            SearchResponse searchResponse = request.slice(sliceBuilder).get();
-            totalResults += searchResponse.getHits().getHits().length;
-            int expectedSliceResults = (int) searchResponse.getHits().getTotalHits().value;
-            int numSliceResults = searchResponse.getHits().getHits().length;
-            String scrollId = searchResponse.getScrollId();
-            for (SearchHit hit : searchResponse.getHits().getHits()) {
-                assertTrue(keys.add(hit.getId()));
-            }
-            while (searchResponse.getHits().getHits().length > 0) {
-                searchResponse = client().prepareSearchScroll("test")
-                    .setScrollId(scrollId)
-                    .setScroll(new Scroll(TimeValue.timeValueSeconds(10)))
-                    .get();
-                scrollId = searchResponse.getScrollId();
+    private void assertSearchSlicesWithScroll(SearchRequestBuilder request, String field, int numSlice, int numDocs) throws Exception {
+        assertBusy(() -> {
+            int totalResults = 0;
+            List<String> keys = new ArrayList<>();
+            for (int id = 0; id < numSlice; id++) {
+                SliceBuilder sliceBuilder = new SliceBuilder(field, id, numSlice);
+                SearchResponse searchResponse = request.slice(sliceBuilder).get();
                 totalResults += searchResponse.getHits().getHits().length;
-                numSliceResults += searchResponse.getHits().getHits().length;
+                int expectedSliceResults = (int) searchResponse.getHits().getTotalHits().value;
+                int numSliceResults = searchResponse.getHits().getHits().length;
+                String scrollId = searchResponse.getScrollId();
                 for (SearchHit hit : searchResponse.getHits().getHits()) {
                     assertTrue(keys.add(hit.getId()));
                 }
+                while (searchResponse.getHits().getHits().length > 0) {
+                    searchResponse = client().prepareSearchScroll("test")
+                        .setScrollId(scrollId)
+                        .setScroll(new Scroll(TimeValue.timeValueSeconds(10)))
+                        .get();
+                    scrollId = searchResponse.getScrollId();
+                    totalResults += searchResponse.getHits().getHits().length;
+                    numSliceResults += searchResponse.getHits().getHits().length;
+                    for (SearchHit hit : searchResponse.getHits().getHits()) {
+                        assertTrue(keys.add(hit.getId()));
+                    }
+                }
+                assertThat(numSliceResults, equalTo(expectedSliceResults));
+                clearScroll(scrollId);
             }
-            assertThat(numSliceResults, equalTo(expectedSliceResults));
-            clearScroll(scrollId);
-        }
-        assertThat(totalResults, equalTo(numDocs));
-        assertThat(keys.size(), equalTo(numDocs));
-        assertThat(new HashSet(keys).size(), equalTo(numDocs));
+            assertThat(totalResults, equalTo(numDocs));
+            assertThat(keys.size(), equalTo(numDocs));
+            assertThat(new HashSet(keys).size(), equalTo(numDocs));
+        }, 30 , TimeUnit.SECONDS);
     }
 
     private Throwable findRootCause(Exception e) {