From 1258401ba8342b1567af7442f95c55df1a302e9b Mon Sep 17 00:00:00 2001
From: Ryan Ernst <ryan@iernst.net>
Date: Tue, 21 Oct 2014 11:57:57 -0700
Subject: [PATCH] Tests: Add forgotten portions of the fix to make the upgrade
 test not do reallcations at the wrong time.

---
 .../rest/action/admin/indices/upgrade/UpgradeTest.java      | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/test/java/org/elasticsearch/rest/action/admin/indices/upgrade/UpgradeTest.java b/src/test/java/org/elasticsearch/rest/action/admin/indices/upgrade/UpgradeTest.java
index 85befd224f1..9bbbd79b1c1 100644
--- a/src/test/java/org/elasticsearch/rest/action/admin/indices/upgrade/UpgradeTest.java
+++ b/src/test/java/org/elasticsearch/rest/action/admin/indices/upgrade/UpgradeTest.java
@@ -41,6 +41,7 @@ import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.index.engine.Segment;
 import org.elasticsearch.node.internal.InternalNode;
 import org.elasticsearch.test.ElasticsearchBackwardsCompatIntegrationTest;
+import org.elasticsearch.test.ElasticsearchIntegrationTest;
 import org.elasticsearch.test.rest.client.http.HttpRequestBuilder;
 import org.elasticsearch.test.rest.client.http.HttpResponse;
 import org.elasticsearch.test.rest.json.JsonPath;
@@ -54,6 +55,7 @@ import java.util.Map;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 
+@ElasticsearchIntegrationTest.ClusterScope(scope = ElasticsearchIntegrationTest.Scope.TEST)   // test scope since we set cluster wide settings
 public class UpgradeTest extends ElasticsearchBackwardsCompatIntegrationTest {
 
     @BeforeClass
@@ -136,6 +138,10 @@ public class UpgradeTest extends ElasticsearchBackwardsCompatIntegrationTest {
         builder.put(EnableAllocationDecider.CLUSTER_ROUTING_ALLOCATION_ENABLE, EnableAllocationDecider.Allocation.NONE);
         client().admin().cluster().prepareUpdateSettings().setTransientSettings(builder).get();
         backwardsCluster().upgradeAllNodes();
+        // we are done - enable allocation again
+        builder = ImmutableSettings.builder();
+        builder.put(EnableAllocationDecider.CLUSTER_ROUTING_ALLOCATION_ENABLE, EnableAllocationDecider.Allocation.ALL);
+        client().admin().cluster().prepareUpdateSettings().setTransientSettings(builder).get();
         ensureGreen();
         logger.debug("--> Nodes upgrade complete");
         logClusterState();