From 3b31f256244626f5f7ef3b3aba74fe778efc4b4e Mon Sep 17 00:00:00 2001
From: Simon Willnauer <simonw@apache.org>
Date: Mon, 2 Jun 2014 11:57:32 +0200
Subject: [PATCH] [TEST] Ensure cluster size reflected in the cluster state

We perform some management operations that require the cluster to be
consistent with respect to the number of nodes in the cluster state
/ visible to the master in order to rely on the ack mechanism. This
only applies to the test infrastructure when nodes are not explicitly
started / stopped as well as while tearing down the cluster and wiping
indices after the tests.
---
 .../elasticsearch/test/ElasticsearchIntegrationTest.java | 5 +++++
 src/test/java/org/elasticsearch/test/TestCluster.java    | 9 +++++++++
 src/test/java/org/elasticsearch/tribe/TribeTests.java    | 2 +-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/test/java/org/elasticsearch/test/ElasticsearchIntegrationTest.java b/src/test/java/org/elasticsearch/test/ElasticsearchIntegrationTest.java
index 337b294fba0..37ba4eaa290 100644
--- a/src/test/java/org/elasticsearch/test/ElasticsearchIntegrationTest.java
+++ b/src/test/java/org/elasticsearch/test/ElasticsearchIntegrationTest.java
@@ -102,6 +102,7 @@ import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
 import static org.elasticsearch.test.TestCluster.clusterName;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoTimeout;
 import static org.hamcrest.Matchers.emptyIterable;
 import static org.hamcrest.Matchers.equalTo;
 
@@ -504,6 +505,7 @@ public abstract class ElasticsearchIntegrationTest extends ElasticsearchTestCase
                     assertThat("test leaves transient cluster metadata behind: " + metaData.transientSettings().getAsMap(), metaData
                             .transientSettings().getAsMap().size(), equalTo(0));
                 }
+                ensureClusterSizeConsistency();
                 immutableCluster().wipe(); // wipe after to make sure we fail in the test that didn't ack the delete
                 immutableCluster().assertAfterTest();
             } finally {
@@ -837,6 +839,9 @@ public abstract class ElasticsearchIntegrationTest extends ElasticsearchTestCase
         return actionGet.getStatus();
     }
 
+    private void ensureClusterSizeConsistency() {
+        assertNoTimeout(client().admin().cluster().prepareHealth().setWaitForNodes(Integer.toString(immutableCluster().size())).get());
+    }
 
     /**
      * Ensures the cluster is in a searchable state for the given indices. This means a searchable copy of each
diff --git a/src/test/java/org/elasticsearch/test/TestCluster.java b/src/test/java/org/elasticsearch/test/TestCluster.java
index 2a10c947da4..33b789c29e3 100644
--- a/src/test/java/org/elasticsearch/test/TestCluster.java
+++ b/src/test/java/org/elasticsearch/test/TestCluster.java
@@ -94,6 +94,7 @@ import static org.apache.lucene.util.LuceneTestCase.rarely;
 import static org.apache.lucene.util.LuceneTestCase.usually;
 import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
 import static org.elasticsearch.node.NodeBuilder.nodeBuilder;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoTimeout;
 
 /**
  * TestCluster manages a set of JVM private nodes and allows convenient access to them.
@@ -397,6 +398,11 @@ public final class TestCluster extends ImmutableTestCluster {
         } catch (Exception e) {
             throw new ElasticsearchException("failed to start nodes", e);
         }
+        if (!futures.isEmpty()) {
+            synchronized (this) {
+                assertNoTimeout(client().admin().cluster().prepareHealth().setWaitForNodes(Integer.toString(nodes.size())).get());
+            }
+        }
     }
 
     /**
@@ -424,6 +430,9 @@ public final class TestCluster extends ImmutableTestCluster {
         for (NodeAndClient toRemove : nodesToRemove) {
             nodes.remove(toRemove.name);
         }
+        if (!nodesToRemove.isEmpty() && size() > 0) {
+            assertNoTimeout(client().admin().cluster().prepareHealth().setWaitForNodes(Integer.toString(nodes.size())).get());
+        }
     }
 
     private NodeAndClient buildNode(Settings settings, Version version) {
diff --git a/src/test/java/org/elasticsearch/tribe/TribeTests.java b/src/test/java/org/elasticsearch/tribe/TribeTests.java
index 769b1848f0b..07f1f775586 100644
--- a/src/test/java/org/elasticsearch/tribe/TribeTests.java
+++ b/src/test/java/org/elasticsearch/tribe/TribeTests.java
@@ -346,4 +346,4 @@ public class TribeTests extends ElasticsearchIntegrationTest {
         }
         return count;
     }
-}
+}
\ No newline at end of file