Add a method that ensures that the cluster is yellow and has no intia…

…lizing shards (#28416)
elastic · Jan 29, 2018 · 43d1dcb · 43d1dcb
1 parent 37c9ac2
commit 43d1dcb
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java
@@ -61,6 +61,7 @@
 import org.elasticsearch.action.index.IndexResponse;
 import org.elasticsearch.action.search.ClearScrollResponse;
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.support.ActiveShardCount;
 import org.elasticsearch.action.support.DefaultShardOperationFailedException;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.client.AdminClient;
@@ -910,17 +911,26 @@ public ClusterHealthStatus ensureGreen(String... indices) {
      * @param timeout time out value to set on {@link org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest}
      */
     public ClusterHealthStatus ensureGreen(TimeValue timeout, String... indices) {
-        return ensureColor(ClusterHealthStatus.GREEN, timeout, indices);
+        return ensureColor(ClusterHealthStatus.GREEN, timeout, false, indices);
     }
 
     /**
      * Ensures the cluster has a yellow state via the cluster health API.
      */
     public ClusterHealthStatus ensureYellow(String... indices) {
-        return ensureColor(ClusterHealthStatus.YELLOW, TimeValue.timeValueSeconds(30), indices);
+        return ensureColor(ClusterHealthStatus.YELLOW, TimeValue.timeValueSeconds(30), false, indices);
     }
 
-    private ClusterHealthStatus ensureColor(ClusterHealthStatus clusterHealthStatus, TimeValue timeout, String... indices) {
+    /**
+     * Ensures the cluster has a yellow state via the cluster health API and ensures the that cluster has no initializing shards
+     * for the given indices
+     */
+    public ClusterHealthStatus ensureYellowAndNoInitializingShards(String... indices) {
+        return ensureColor(ClusterHealthStatus.YELLOW, TimeValue.timeValueSeconds(30), true, indices);
+    }
+
+    private ClusterHealthStatus ensureColor(ClusterHealthStatus clusterHealthStatus, TimeValue timeout, boolean waitForNoInitializingShards,
+                                            String... indices) {
         String color = clusterHealthStatus.name().toLowerCase(Locale.ROOT);
         String method = "ensure" + Strings.capitalize(color);
 
@@ -929,6 +939,7 @@ private ClusterHealthStatus ensureColor(ClusterHealthStatus clusterHealthStatus,
             .waitForStatus(clusterHealthStatus)
             .waitForEvents(Priority.LANGUID)
             .waitForNoRelocatingShards(true)
+            .waitForNoInitializingShards(waitForNoInitializingShards)
             // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
             // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
             // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently