From 6e5438b236983cdf763043830d7bc48ab607d5e5 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Tue, 1 Oct 2024 05:44:26 +0000
Subject: [PATCH] Add assertBusy to failing snapshot tests (#16146)

Signed-off-by: Lakshya Taragi <lakshya.taragi@gmail.com>
(cherry picked from commit 43e7597cdc5ba2c1852ec1796628f948633f0c57)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../snapshots/SnapshotStatusApisIT.java       | 40 ++++++++++++-------
 1 file changed, 25 insertions(+), 15 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/snapshots/SnapshotStatusApisIT.java b/server/src/internalClusterTest/java/org/opensearch/snapshots/SnapshotStatusApisIT.java
index c3214022df663..8b6869aa1d81a 100644
--- a/server/src/internalClusterTest/java/org/opensearch/snapshots/SnapshotStatusApisIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/snapshots/SnapshotStatusApisIT.java
@@ -116,7 +116,7 @@ public void testStatusApiConsistency() {
         assertEquals(snapshotStatus.getStats().getTime(), snapshotInfo.endTime() - snapshotInfo.startTime());
     }
 
-    public void testStatusAPICallForShallowCopySnapshot() {
+    public void testStatusAPICallForShallowCopySnapshot() throws Exception {
         disableRepoConsistencyCheck("Remote store repository is being used for the test");
         internalCluster().startClusterManagerOnlyNode();
         internalCluster().startDataOnlyNode();
@@ -136,15 +136,24 @@ public void testStatusAPICallForShallowCopySnapshot() {
         final String snapshot = "snapshot";
         createFullSnapshot(snapshotRepoName, snapshot);
 
-        final SnapshotStatus snapshotStatus = getSnapshotStatus(snapshotRepoName, snapshot);
-        assertThat(snapshotStatus.getState(), is(SnapshotsInProgress.State.SUCCESS));
+        assertBusy(() -> {
+            final SnapshotStatus snapshotStatus = client().admin()
+                .cluster()
+                .prepareSnapshotStatus(snapshotRepoName)
+                .setSnapshots(snapshot)
+                .execute()
+                .actionGet()
+                .getSnapshots()
+                .get(0);
+            assertThat(snapshotStatus.getState(), is(SnapshotsInProgress.State.SUCCESS));
 
-        final SnapshotIndexShardStatus snapshotShardState = stateFirstShard(snapshotStatus, indexName);
-        assertThat(snapshotShardState.getStage(), is(SnapshotIndexShardStage.DONE));
-        assertThat(snapshotShardState.getStats().getTotalFileCount(), greaterThan(0));
-        assertThat(snapshotShardState.getStats().getTotalSize(), greaterThan(0L));
-        assertThat(snapshotShardState.getStats().getIncrementalFileCount(), greaterThan(0));
-        assertThat(snapshotShardState.getStats().getIncrementalSize(), greaterThan(0L));
+            final SnapshotIndexShardStatus snapshotShardState = stateFirstShard(snapshotStatus, indexName);
+            assertThat(snapshotShardState.getStage(), is(SnapshotIndexShardStage.DONE));
+            assertThat(snapshotShardState.getStats().getTotalFileCount(), greaterThan(0));
+            assertThat(snapshotShardState.getStats().getTotalSize(), greaterThan(0L));
+            assertThat(snapshotShardState.getStats().getIncrementalFileCount(), greaterThan(0));
+            assertThat(snapshotShardState.getStats().getIncrementalSize(), greaterThan(0L));
+        }, 20, TimeUnit.SECONDS);
     }
 
     public void testStatusAPICallInProgressSnapshot() throws Exception {
@@ -193,7 +202,7 @@ public void testExceptionOnMissingSnapBlob() throws IOException {
         );
     }
 
-    public void testExceptionOnMissingShardLevelSnapBlob() throws IOException {
+    public void testExceptionOnMissingShardLevelSnapBlob() throws Exception {
         disableRepoConsistencyCheck("This test intentionally corrupts the repository");
 
         final Path repoPath = randomRepoPath();
@@ -216,11 +225,12 @@ public void testExceptionOnMissingShardLevelSnapBlob() throws IOException {
             repoPath.resolve(resolvePath(indexId, "0"))
                 .resolve(BlobStoreRepository.SNAPSHOT_PREFIX + snapshotInfo.snapshotId().getUUID() + ".dat")
         );
-
-        expectThrows(
-            SnapshotMissingException.class,
-            () -> client().admin().cluster().prepareSnapshotStatus("test-repo").setSnapshots("test-snap").execute().actionGet()
-        );
+        assertBusy(() -> {
+            expectThrows(
+                SnapshotMissingException.class,
+                () -> client().admin().cluster().prepareSnapshotStatus("test-repo").setSnapshots("test-snap").execute().actionGet()
+            );
+        }, 20, TimeUnit.SECONDS);
     }
 
     public void testGetSnapshotsWithoutIndices() throws Exception {