6161import org .elasticsearch .action .index .IndexResponse ;
6262import org .elasticsearch .action .search .ClearScrollResponse ;
6363import org .elasticsearch .action .search .SearchResponse ;
64+ import org .elasticsearch .action .support .ActiveShardCount ;
6465import org .elasticsearch .action .support .DefaultShardOperationFailedException ;
6566import org .elasticsearch .action .support .IndicesOptions ;
6667import org .elasticsearch .client .AdminClient ;
@@ -910,17 +911,26 @@ public ClusterHealthStatus ensureGreen(String... indices) {
910911 * @param timeout time out value to set on {@link org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest}
911912 */
912913 public ClusterHealthStatus ensureGreen (TimeValue timeout , String ... indices ) {
913- return ensureColor (ClusterHealthStatus .GREEN , timeout , indices );
914+ return ensureColor (ClusterHealthStatus .GREEN , timeout , false , indices );
914915 }
915916
916917 /**
917918 * Ensures the cluster has a yellow state via the cluster health API.
918919 */
919920 public ClusterHealthStatus ensureYellow (String ... indices ) {
920- return ensureColor (ClusterHealthStatus .YELLOW , TimeValue .timeValueSeconds (30 ), indices );
921+ return ensureColor (ClusterHealthStatus .YELLOW , TimeValue .timeValueSeconds (30 ), false , indices );
921922 }
922923
923- private ClusterHealthStatus ensureColor (ClusterHealthStatus clusterHealthStatus , TimeValue timeout , String ... indices ) {
924+ /**
925+ * Ensures the cluster has a yellow state via the cluster health API and ensures the that cluster has no initializing shards
926+ * for the given indices
927+ */
928+ public ClusterHealthStatus ensureYellowAndNoInitializingShards (String ... indices ) {
929+ return ensureColor (ClusterHealthStatus .YELLOW , TimeValue .timeValueSeconds (30 ), true , indices );
930+ }
931+
932+ private ClusterHealthStatus ensureColor (ClusterHealthStatus clusterHealthStatus , TimeValue timeout , boolean waitForNoInitializingShards ,
933+ String ... indices ) {
924934 String color = clusterHealthStatus .name ().toLowerCase (Locale .ROOT );
925935 String method = "ensure" + Strings .capitalize (color );
926936
@@ -929,6 +939,7 @@ private ClusterHealthStatus ensureColor(ClusterHealthStatus clusterHealthStatus,
929939 .waitForStatus (clusterHealthStatus )
930940 .waitForEvents (Priority .LANGUID )
931941 .waitForNoRelocatingShards (true )
942+ .waitForNoInitializingShards (waitForNoInitializingShards )
932943 // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
933944 // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
934945 // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
0 commit comments