do not use sleep

Author: dnhatn

core/src/test/java/org/elasticsearch/snapshots/SnapshotShardsServiceIT.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeUnit;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.snapshots.mockstore.MockRepository;
 import org.elasticsearch.test.ESIntegTestCase;
@@ -31,10 +32,14 @@
 
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.everyItem;
+import static org.hamcrest.Matchers.hasSize;
 
 @ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST, numDataNodes = 0, transportClientRatio = 0)
 public class SnapshotShardsServiceIT extends AbstractSnapshotIntegTestCase {
@@ -71,6 +76,9 @@ public void testRetryPostingSnapshotStatusMessages() throws Exception {
             .get();
         waitForBlock(blockedNode, "test-repo", TimeValue.timeValueSeconds(60));
 
+        final SnapshotId snapshotId = client().admin().cluster().prepareGetSnapshots("test-repo").setSnapshots("test-snap")
+            .get().getSnapshots().get(0).snapshotId();
+
         logger.info("--> start disrupting cluster");
         final NetworkDisruption networkDisruption = new NetworkDisruption(new NetworkDisruption.TwoPartitions(masterNode, dataNode),
             NetworkDisruption.NetworkDelay.random(random()));
@@ -79,16 +87,27 @@ public void testRetryPostingSnapshotStatusMessages() throws Exception {
 
         logger.info("--> unblocking repository");
         unblockNode("test-repo", blockedNode);
-        Thread.sleep(200);
+
+        // Retrieve snapshot status from the data node.
+        SnapshotShardsService snapshotShardsService = internalCluster().getInstance(SnapshotShardsService.class, blockedNode);
+        assertBusy(() -> {
+            final Snapshot snapshot = new Snapshot("test-repo", snapshotId);
+            List<IndexShardSnapshotStatus.Stage> stages = snapshotShardsService.currentSnapshotShards(snapshot)
+                .values().stream().map(IndexShardSnapshotStatus::stage).collect(Collectors.toList());
+            assertThat(stages, hasSize(shards));
+            assertThat(stages, everyItem(equalTo(IndexShardSnapshotStatus.Stage.DONE)));
+        });
+
         logger.info("--> stop disrupting cluster");
+        networkDisruption.stopDisrupting();
         internalCluster().clearDisruptionScheme(true);
 
         assertBusy(() -> {
             GetSnapshotsResponse snapshotsStatusResponse = client().admin().cluster()
                 .prepareGetSnapshots("test-repo")
                 .setSnapshots("test-snap").get();
-            logger.info("Status size [{}]", snapshotsStatusResponse.getSnapshots().get(0).successfulShards());
             SnapshotInfo snapshotInfo = snapshotsStatusResponse.getSnapshots().get(0);
+            logger.info("Snapshot status [{}], successfulShards [{}]", snapshotInfo.state(), snapshotInfo.successfulShards());
             assertThat(snapshotInfo.state(), equalTo(SnapshotState.SUCCESS));
             assertThat(snapshotInfo.successfulShards(), equalTo(shards));
         }, 10, TimeUnit.SECONDS);