From e297b21c2ce6f887919d1d06ed76b8d6c9d27a13 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Wed, 1 Nov 2023 13:27:57 -0500
Subject: [PATCH] Fix testDeleteAllAndListAllPits and testDeleteWhileSearch
 flaky tests (#10946) (#11049)

(cherry picked from commit 8673fa937db405b8d614f8d4a02c0aa52587c037)

Signed-off-by: Bharathwaj G <bharath78910@gmail.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../java/org/opensearch/client/PitIT.java     | 21 +++++++++----------
 .../search/pit/DeletePitMultiNodeIT.java      |  6 +++++-
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/client/rest-high-level/src/test/java/org/opensearch/client/PitIT.java b/client/rest-high-level/src/test/java/org/opensearch/client/PitIT.java
index 5ec1da77a6795..6e4b83d2420fa 100644
--- a/client/rest-high-level/src/test/java/org/opensearch/client/PitIT.java
+++ b/client/rest-high-level/src/test/java/org/opensearch/client/PitIT.java
@@ -24,7 +24,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
@@ -72,7 +71,7 @@ public void testCreateAndDeletePit() throws IOException {
         assertTrue(deletePitResponse.getDeletePitResults().get(0).getPitId().equals(createPitResponse.getId()));
     }
 
-    public void testDeleteAllAndListAllPits() throws IOException, InterruptedException {
+    public void testDeleteAllAndListAllPits() throws Exception {
         CreatePitRequest pitRequest = new CreatePitRequest(new TimeValue(1, TimeUnit.DAYS), true, "index");
         CreatePitResponse pitResponse = execute(pitRequest, highLevelClient()::createPit, highLevelClient()::createPitAsync);
         CreatePitResponse pitResponse1 = execute(pitRequest, highLevelClient()::createPit, highLevelClient()::createPitAsync);
@@ -91,11 +90,9 @@ public void testDeleteAllAndListAllPits() throws IOException, InterruptedExcepti
         List<String> pits = getAllPitResponse.getPitInfos().stream().map(r -> r.getPitId()).collect(Collectors.toList());
         assertTrue(pits.contains(pitResponse.getId()));
         assertTrue(pits.contains(pitResponse1.getId()));
-        CountDownLatch countDownLatch = new CountDownLatch(1);
         ActionListener<DeletePitResponse> deletePitListener = new ActionListener<>() {
             @Override
             public void onResponse(DeletePitResponse response) {
-                countDownLatch.countDown();
                 for (DeletePitInfo deletePitInfo : response.getDeletePitResults()) {
                     assertTrue(deletePitInfo.isSuccessful());
                 }
@@ -103,19 +100,20 @@ public void onResponse(DeletePitResponse response) {
 
             @Override
             public void onFailure(Exception e) {
-                countDownLatch.countDown();
                 if (!(e instanceof OpenSearchStatusException)) {
                     throw new AssertionError("Delete all failed");
                 }
             }
         };
         final CreatePitResponse pitResponse3 = execute(pitRequest, highLevelClient()::createPit, highLevelClient()::createPitAsync);
-
+        assertTrue(pitResponse3.getId() != null);
         ActionListener<GetAllPitNodesResponse> getPitsListener = new ActionListener<GetAllPitNodesResponse>() {
             @Override
             public void onResponse(GetAllPitNodesResponse response) {
                 List<String> pits = response.getPitInfos().stream().map(r -> r.getPitId()).collect(Collectors.toList());
                 assertTrue(pits.contains(pitResponse3.getId()));
+                // delete all pits
+                highLevelClient().deleteAllPitsAsync(RequestOptions.DEFAULT, deletePitListener);
             }
 
             @Override
@@ -126,11 +124,12 @@ public void onFailure(Exception e) {
             }
         };
         highLevelClient().getAllPitsAsync(RequestOptions.DEFAULT, getPitsListener);
-        highLevelClient().deleteAllPitsAsync(RequestOptions.DEFAULT, deletePitListener);
-        assertTrue(countDownLatch.await(10, TimeUnit.SECONDS));
+
         // validate no pits case
-        getAllPitResponse = highLevelClient().getAllPits(RequestOptions.DEFAULT);
-        assertTrue(getAllPitResponse.getPitInfos().size() == 0);
-        highLevelClient().deleteAllPitsAsync(RequestOptions.DEFAULT, deletePitListener);
+        assertBusy(() -> {
+            GetAllPitNodesResponse getAllPitResponse1 = highLevelClient().getAllPits(RequestOptions.DEFAULT);
+            assertTrue(getAllPitResponse1.getPitInfos().size() == 0);
+            highLevelClient().deleteAllPitsAsync(RequestOptions.DEFAULT, deletePitListener);
+        });
     }
 }
diff --git a/server/src/internalClusterTest/java/org/opensearch/search/pit/DeletePitMultiNodeIT.java b/server/src/internalClusterTest/java/org/opensearch/search/pit/DeletePitMultiNodeIT.java
index 4a178e7066846..8ae652082f653 100644
--- a/server/src/internalClusterTest/java/org/opensearch/search/pit/DeletePitMultiNodeIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/search/pit/DeletePitMultiNodeIT.java
@@ -309,7 +309,11 @@ public void testDeleteWhileSearch() throws Exception {
     private void verifySearchContextMissingException(ShardSearchFailure[] failures) {
         for (ShardSearchFailure failure : failures) {
             Throwable cause = ExceptionsHelper.unwrapCause(failure.getCause());
-            assertTrue(failure.toString(), cause instanceof SearchContextMissingException);
+            if (failure.toString().contains("reader_context is already closed can't increment refCount current count")) {
+                // this is fine, expected search error when context is already deleted
+            } else {
+                assertTrue(failure.toString(), cause instanceof SearchContextMissingException);
+            }
         }
     }