Dry up Snapshot Integ Tests some more

server/src/internalClusterTest/java/org/elasticsearch/repositories/blobstore/BlobStoreRepositoryCleanupIT.java

@@ -31,7 +31,6 @@
 
 import java.io.ByteArrayInputStream;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertFutureThrows;
 import static org.hamcrest.Matchers.is;
@@ -49,11 +48,8 @@ public void testMasterFailoverDuringCleanup() throws Exception {
         ensureStableCluster(nodeCount - 1);
 
         logger.info("-->  wait for cleanup to finish and disappear from cluster state");
-        assertBusy(() -> {
-            RepositoryCleanupInProgress cleanupInProgress =
-                client().admin().cluster().prepareState().get().getState().custom(RepositoryCleanupInProgress.TYPE);
-            assertFalse(cleanupInProgress.hasCleanupInProgress());
-        }, 30, TimeUnit.SECONDS);
+        awaitClusterState(state ->
+                state.custom(RepositoryCleanupInProgress.TYPE, RepositoryCleanupInProgress.EMPTY).hasCleanupInProgress() == false);
     }
 
     public void testRepeatCleanupsDontRemove() throws Exception {
@@ -71,11 +67,8 @@ public void testRepeatCleanupsDontRemove() throws Exception {
         unblockNode("test-repo", masterNode);
 
         logger.info("-->  wait for cleanup to finish and disappear from cluster state");
-        assertBusy(() -> {
-            RepositoryCleanupInProgress cleanupInProgress =
-                client().admin().cluster().prepareState().get().getState().custom(RepositoryCleanupInProgress.TYPE);
-            assertFalse(cleanupInProgress.hasCleanupInProgress());
-        }, 30, TimeUnit.SECONDS);
+        awaitClusterState(state ->
+                state.custom(RepositoryCleanupInProgress.TYPE, RepositoryCleanupInProgress.EMPTY).hasCleanupInProgress() == false);
     }
 
     private String startBlockedCleanup(String repoName) throws Exception {

server/src/internalClusterTest/java/org/elasticsearch/snapshots/CorruptedBlobStoreRepositoryIT.java

@@ -24,10 +24,8 @@
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.ClusterState;
-import org.elasticsearch.cluster.ClusterStateUpdateTask;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.metadata.RepositoriesMetadata;
-import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeUnit;
@@ -149,29 +147,10 @@ public void testFindDanglingLatestGeneration() throws Exception {
         Files.move(repo.resolve("index-" + beforeMoveGen), repo.resolve("index-" + (beforeMoveGen + 1)));
 
         logger.info("--> set next generation as pending in the cluster state");
-        final PlainActionFuture<Void> csUpdateFuture = PlainActionFuture.newFuture();
-        internalCluster().getCurrentMasterNodeInstance(ClusterService.class).submitStateUpdateTask("set pending generation",
-            new ClusterStateUpdateTask() {
-                @Override
-                public ClusterState execute(ClusterState currentState) {
-                    return ClusterState.builder(currentState).metadata(Metadata.builder(currentState.getMetadata())
-                        .putCustom(RepositoriesMetadata.TYPE,
-                            currentState.metadata().<RepositoriesMetadata>custom(RepositoriesMetadata.TYPE).withUpdatedGeneration(
-                                repository.getMetadata().name(), beforeMoveGen, beforeMoveGen + 1)).build()).build();
-                }
-
-                @Override
-                public void onFailure(String source, Exception e) {
-                    csUpdateFuture.onFailure(e);
-                }
-
-                @Override
-                public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) {
-                    csUpdateFuture.onResponse(null);
-                }
-            }
-        );
-        csUpdateFuture.get();
+        updateClusterState(currentState -> ClusterState.builder(currentState).metadata(Metadata.builder(currentState.getMetadata())
+                .putCustom(RepositoriesMetadata.TYPE,
+                        currentState.metadata().<RepositoriesMetadata>custom(RepositoriesMetadata.TYPE).withUpdatedGeneration(
+                                repository.getMetadata().name(), beforeMoveGen, beforeMoveGen + 1)).build()).build());
 
         logger.info("--> full cluster restart");
         internalCluster().fullRestart();

server/src/internalClusterTest/java/org/elasticsearch/snapshots/SnapshotShardsServiceIT.java

@@ -19,7 +19,6 @@
 
 package org.elasticsearch.snapshots;
 
-import org.elasticsearch.action.admin.cluster.snapshots.get.GetSnapshotsResponse;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
 import org.elasticsearch.plugins.Plugin;
@@ -66,8 +65,7 @@ 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("test-repo").get(0).snapshotId();
+        final SnapshotId snapshotId = getSnapshot("test-repo", "test-snap").snapshotId();
 
         logger.info("--> start disrupting cluster");
         final NetworkDisruption networkDisruption = isolateMasterDisruption(NetworkDisruption.NetworkDelay.random(random()));
@@ -92,10 +90,7 @@ public void testRetryPostingSnapshotStatusMessages() throws Exception {
         internalCluster().clearDisruptionScheme(true);
 
         assertBusy(() -> {
-            GetSnapshotsResponse snapshotsStatusResponse = client().admin().cluster()
-                .prepareGetSnapshots("test-repo")
-                .setSnapshots("test-snap").get();
-            SnapshotInfo snapshotInfo = snapshotsStatusResponse.getSnapshots("test-repo").get(0);
+            SnapshotInfo snapshotInfo = getSnapshot("test-repo", "test-snap");
             logger.info("Snapshot status [{}], successfulShards [{}]", snapshotInfo.state(), snapshotInfo.successfulShards());
             assertThat(snapshotInfo.state(), equalTo(SnapshotState.SUCCESS));
             assertThat(snapshotInfo.successfulShards(), equalTo(shards));

server/src/internalClusterTest/java/org/elasticsearch/snapshots/SnapshotStatusApisIT.java

@@ -28,7 +28,6 @@
 import org.elasticsearch.action.admin.cluster.snapshots.status.SnapshotStats;
 import org.elasticsearch.action.admin.cluster.snapshots.status.SnapshotStatus;
 import org.elasticsearch.action.admin.cluster.snapshots.status.SnapshotsStatusRequest;
-import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.SnapshotsInProgress;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
@@ -49,8 +48,6 @@
 public class SnapshotStatusApisIT extends AbstractSnapshotIntegTestCase {
 
     public void testStatusApiConsistency() {
-        Client client = client();
-
         createRepository("test-repo", "fs");
 
         createIndex("test-idx-1", "test-idx-2", "test-idx-3");
@@ -66,14 +63,13 @@ public void testStatusApiConsistency() {
 
         createFullSnapshot("test-repo", "test-snap");
 
-        List<SnapshotInfo> snapshotInfos =
-            client.admin().cluster().prepareGetSnapshots("test-repo").get().getSnapshots("test-repo");
+        List<SnapshotInfo> snapshotInfos = clusterAdmin().prepareGetSnapshots("test-repo").get().getSnapshots("test-repo");
         assertThat(snapshotInfos.size(), equalTo(1));
         SnapshotInfo snapshotInfo = snapshotInfos.get(0);
         assertThat(snapshotInfo.state(), equalTo(SnapshotState.SUCCESS));
         assertThat(snapshotInfo.version(), equalTo(Version.CURRENT));
 
-        final List<SnapshotStatus> snapshotStatus = client.admin().cluster().snapshotsStatus(
+        final List<SnapshotStatus> snapshotStatus = clusterAdmin().snapshotsStatus(
             new SnapshotsStatusRequest("test-repo", new String[]{"test-snap"})).actionGet().getSnapshots();
         assertThat(snapshotStatus.size(), equalTo(1));
         final SnapshotStatus snStatus = snapshotStatus.get(0);
@@ -82,8 +78,6 @@ public void testStatusApiConsistency() {
     }
 
     public void testStatusAPICallInProgressSnapshot() throws Exception {
-        Client client = client();
-
         createRepository("test-repo", "mock", Settings.builder().put("location", randomRepoPath()).put("block_on_data", true));
 
         createIndex("test-idx-1");
@@ -96,13 +90,12 @@ public void testStatusAPICallInProgressSnapshot() throws Exception {
         refresh();
 
         logger.info("--> snapshot");
-        ActionFuture<CreateSnapshotResponse> createSnapshotResponseActionFuture =
-            client.admin().cluster().prepareCreateSnapshot("test-repo", "test-snap").setWaitForCompletion(true).execute();
+        ActionFuture<CreateSnapshotResponse> createSnapshotResponseActionFuture = startFullSnapshot("test-repo", "test-snap");
 
         logger.info("--> wait for data nodes to get blocked");
         waitForBlockOnAnyDataNode("test-repo", TimeValue.timeValueMinutes(1));
         awaitNumberOfSnapshotsInProgress(1);
-        assertEquals(SnapshotsInProgress.State.STARTED, client.admin().cluster().prepareSnapshotStatus("test-repo")
+        assertEquals(SnapshotsInProgress.State.STARTED, client().admin().cluster().prepareSnapshotStatus("test-repo")
                 .setSnapshots("test-snap").get().getSnapshots().get(0).getState());
 
         logger.info("--> unblock all data nodes");

test/framework/src/main/java/org/elasticsearch/snapshots/AbstractSnapshotIntegTestCase.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateObserver;
+import org.elasticsearch.cluster.ClusterStateUpdateTask;
 import org.elasticsearch.cluster.SnapshotDeletionsInProgress;
 import org.elasticsearch.cluster.SnapshotsInProgress;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
@@ -82,7 +83,7 @@
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.Matchers.empty;
-import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
 
 public abstract class AbstractSnapshotIntegTestCase extends ESIntegTestCase {
@@ -122,11 +123,11 @@ public void verifyNoLeakedListeners() throws Exception {
     @After
     public void assertRepoConsistency() {
         if (skipRepoConsistencyCheckReason == null) {
-            client().admin().cluster().prepareGetRepositories().get().repositories().forEach(repositoryMetadata -> {
+            clusterAdmin().prepareGetRepositories().get().repositories().forEach(repositoryMetadata -> {
                 final String name = repositoryMetadata.name();
                 if (repositoryMetadata.settings().getAsBoolean("readonly", false) == false) {
-                    client().admin().cluster().prepareDeleteSnapshot(name, OLD_VERSION_SNAPSHOT_PREFIX + "*").get();
-                    client().admin().cluster().prepareCleanupRepository(name).get();
+                    clusterAdmin().prepareDeleteSnapshot(name, OLD_VERSION_SNAPSHOT_PREFIX + "*").get();
+                    clusterAdmin().prepareCleanupRepository(name).get();
                 }
                 BlobStoreTestUtil.assertRepoConsistency(internalCluster(), name);
             });
@@ -202,12 +203,10 @@ public static void waitForBlock(String node, String repository, TimeValue timeou
     public SnapshotInfo waitForCompletion(String repository, String snapshotName, TimeValue timeout) throws InterruptedException {
         long start = System.currentTimeMillis();
         while (System.currentTimeMillis() - start < timeout.millis()) {
-            List<SnapshotInfo> snapshotInfos = client().admin().cluster().prepareGetSnapshots(repository).setSnapshots(snapshotName)
-                .get().getSnapshots(repository);
-            assertThat(snapshotInfos.size(), equalTo(1));
-            if (snapshotInfos.get(0).state().completed()) {
+            final SnapshotInfo snapshotInfo = getSnapshot(repository, snapshotName);
+            if (snapshotInfo.state().completed()) {
                 // Make sure that snapshot clean up operations are finished
-                ClusterStateResponse stateResponse = client().admin().cluster().prepareState().get();
+                ClusterStateResponse stateResponse = clusterAdmin().prepareState().get();
                 boolean found = false;
                 for (SnapshotsInProgress.Entry entry :
                     stateResponse.getState().custom(SnapshotsInProgress.TYPE, SnapshotsInProgress.EMPTY).entries()) {
@@ -218,7 +217,7 @@ public SnapshotInfo waitForCompletion(String repository, String snapshotName, Ti
                     }
                 }
                 if (found == false) {
-                    return snapshotInfos.get(0);
+                    return snapshotInfo;
                 }
             }
             Thread.sleep(100);
@@ -307,7 +306,7 @@ public void unblockNode(final String repository, final String node) {
 
     protected void createRepository(String repoName, String type, Settings.Builder settings) {
         logger.info("--> creating repository [{}] [{}]", repoName, type);
-        assertAcked(client().admin().cluster().preparePutRepository(repoName)
+        assertAcked(clusterAdmin().preparePutRepository(repoName)
             .setType(type).setSettings(settings));
     }
 
@@ -349,7 +348,7 @@ protected void maybeInitWithOldSnapshotVersion(String repoName, Path repoPath) t
     protected String initWithSnapshotVersion(String repoName, Path repoPath, Version version) throws IOException {
         assertThat("This hack only works on an empty repository", getRepositoryData(repoName).getSnapshotIds(), empty());
         final String oldVersionSnapshot = OLD_VERSION_SNAPSHOT_PREFIX + version.id;
-        final CreateSnapshotResponse createSnapshotResponse = client().admin().cluster()
+        final CreateSnapshotResponse createSnapshotResponse = clusterAdmin()
                 .prepareCreateSnapshot(repoName, oldVersionSnapshot).setIndices("does-not-exist-for-sure-*")
                 .setWaitForCompletion(true).get();
         assertThat(createSnapshotResponse.getSnapshotInfo().totalShards(), is(0));
@@ -372,7 +371,7 @@ protected String initWithSnapshotVersion(String repoName, Path repoPath, Version
 
     protected SnapshotInfo createFullSnapshot(String repoName, String snapshotName) {
         logger.info("--> creating full snapshot [{}] in [{}]", snapshotName, repoName);
-        CreateSnapshotResponse createSnapshotResponse = client().admin().cluster().prepareCreateSnapshot(repoName, snapshotName)
+        CreateSnapshotResponse createSnapshotResponse = clusterAdmin().prepareCreateSnapshot(repoName, snapshotName)
             .setIncludeGlobalState(true)
             .setWaitForCompletion(true)
             .get();
@@ -416,7 +415,7 @@ protected void assertDocCount(String index, long count) {
      * @param metadata     snapshot metadata to write (as returned by {@link SnapshotInfo#userMetadata()})
      */
     protected void addBwCFailedSnapshot(String repoName, String snapshotName, Map<String, Object> metadata) throws Exception {
-        final ClusterState state = client().admin().cluster().prepareState().get().getState();
+        final ClusterState state = clusterAdmin().prepareState().get().getState();
         final RepositoriesMetadata repositoriesMetadata = state.metadata().custom(RepositoriesMetadata.TYPE);
         assertNotNull(repositoriesMetadata);
         final RepositoryMetadata initialRepoMetadata = repositoriesMetadata.repository(repoName);
@@ -486,7 +485,7 @@ protected ActionFuture<CreateSnapshotResponse> startFullSnapshot(String repoName
 
     protected ActionFuture<CreateSnapshotResponse> startFullSnapshot(String repoName, String snapshotName, boolean partial) {
         logger.info("--> creating full snapshot [{}] to repo [{}]", snapshotName, repoName);
-        return client().admin().cluster().prepareCreateSnapshot(repoName, snapshotName).setWaitForCompletion(true)
+        return clusterAdmin().prepareCreateSnapshot(repoName, snapshotName).setWaitForCompletion(true)
                 .setPartial(partial).execute();
     }
 
@@ -517,6 +516,35 @@ protected void createIndexWithContent(String indexName, Settings indexSettings)
 
     protected ActionFuture<AcknowledgedResponse> startDeleteSnapshot(String repoName, String snapshotName) {
         logger.info("--> deleting snapshot [{}] from repo [{}]", snapshotName, repoName);
-        return client().admin().cluster().prepareDeleteSnapshot(repoName, snapshotName).execute();
+        return clusterAdmin().prepareDeleteSnapshot(repoName, snapshotName).execute();
+    }
+
+    protected void updateClusterState(final Function<ClusterState, ClusterState> updater) throws Exception {
+        final PlainActionFuture<Void> future = PlainActionFuture.newFuture();
+        final ClusterService clusterService = internalCluster().getCurrentMasterNodeInstance(ClusterService.class);
+        clusterService.submitStateUpdateTask("test", new ClusterStateUpdateTask() {
+            @Override
+            public ClusterState execute(ClusterState currentState) {
+                return updater.apply(currentState);
+            }
+
+            @Override
+            public void onFailure(String source, Exception e) {
+                future.onFailure(e);
+            }
+
+            @Override
+            public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) {
+                future.onResponse(null);
+            }
+        });
+        future.get();
+    }
+
+    protected SnapshotInfo getSnapshot(String repository, String snapshot) {
+        final List<SnapshotInfo> snapshotInfos = clusterAdmin().prepareGetSnapshots(repository).setSnapshots(snapshot)
+                .get().getSnapshots(repository);
+        assertThat(snapshotInfos, hasSize(1));
+        return snapshotInfos.get(0);
     }
 }

test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java

@@ -58,6 +58,7 @@
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.client.AdminClient;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.client.ClusterAdminClient;
 import org.elasticsearch.client.Requests;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.client.RestClientBuilder;
@@ -1329,6 +1330,13 @@ protected AdminClient admin() {
         return client().admin();
     }
 
+    /**
+     * Returns a random cluster admin client. This client can be pointing to any of the nodes in the cluster.
+     */
+    protected ClusterAdminClient clusterAdmin() {
+        return admin().cluster();
+    }
+
     /**
      * Convenience method that forwards to {@link #indexRandom(boolean, List)}.
      */