From 3d49b4ed3a1bb083387fbf4b5c3d9dea81c6e1ae Mon Sep 17 00:00:00 2001
From: Shay Banon <kimchy@gmail.com>
Date: Mon, 3 Oct 2011 13:48:10 +0200
Subject: [PATCH] add a simple integration level test for awareness

---
 .../allocation/AwarenessAllocationTests.java  | 88 +++++++++++++++++++
 1 file changed, 88 insertions(+)
 create mode 100644 modules/test/integration/src/test/java/org/elasticsearch/test/integration/cluster/allocation/AwarenessAllocationTests.java

diff --git a/modules/test/integration/src/test/java/org/elasticsearch/test/integration/cluster/allocation/AwarenessAllocationTests.java b/modules/test/integration/src/test/java/org/elasticsearch/test/integration/cluster/allocation/AwarenessAllocationTests.java
new file mode 100644
index 00000000000..ca911c0010f
--- /dev/null
+++ b/modules/test/integration/src/test/java/org/elasticsearch/test/integration/cluster/allocation/AwarenessAllocationTests.java
@@ -0,0 +1,88 @@
+/*
+ * Licensed to Elastic Search and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Elastic Search licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.test.integration.cluster.allocation;
+
+import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.routing.IndexRoutingTable;
+import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.logging.ESLogger;
+import org.elasticsearch.common.logging.Loggers;
+import org.elasticsearch.common.settings.ImmutableSettings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.trove.map.hash.TObjectIntHashMap;
+import org.elasticsearch.test.integration.AbstractNodesTests;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+import static org.hamcrest.MatcherAssert.*;
+import static org.hamcrest.Matchers.*;
+
+/**
+ */
+@Test
+public class AwarenessAllocationTests extends AbstractNodesTests {
+
+    private final ESLogger logger = Loggers.getLogger(AwarenessAllocationTests.class);
+
+    @AfterMethod public void cleanAndCloseNodes() throws Exception {
+        closeAllNodes();
+    }
+
+    @Test public void testSimpleAwareness() throws Exception {
+        Settings commonSettings = ImmutableSettings.settingsBuilder()
+                .put("cluster.routing.schedule", "10ms")
+                .put("cluster.routing.allocation.awareness.attributes", "rack_id")
+                .build();
+
+
+        logger.info("--> starting 2 nodes on the same rack");
+        startNode("node1", ImmutableSettings.settingsBuilder().put(commonSettings).put("node.rack_id", "rack_1"));
+        startNode("node2", ImmutableSettings.settingsBuilder().put(commonSettings).put("node.rack_id", "rack_1"));
+
+        client("node1").admin().indices().prepareCreate("test1").execute().actionGet();
+        client("node1").admin().indices().prepareCreate("test2").execute().actionGet();
+
+        ClusterHealthResponse health = client("node1").admin().cluster().prepareHealth().setWaitForGreenStatus().execute().actionGet();
+        assertThat(health.timedOut(), equalTo(false));
+
+        logger.info("--> starting 1 node on a different rack");
+        startNode("node3", ImmutableSettings.settingsBuilder().put(commonSettings).put("node.rack_id", "rack_2"));
+
+        Thread.sleep(500);
+
+        health = client("node1").admin().cluster().prepareHealth().setWaitForGreenStatus().setWaitForNodes("3").setWaitForRelocatingShards(0).execute().actionGet();
+        assertThat(health.timedOut(), equalTo(false));
+
+        ClusterState clusterState = client("node1").admin().cluster().prepareState().execute().actionGet().state();
+        //System.out.println(clusterState.routingTable().prettyPrint());
+        // verify that we have 10 shards on node3
+        TObjectIntHashMap<String> counts = new TObjectIntHashMap<String>();
+        for (IndexRoutingTable indexRoutingTable : clusterState.routingTable()) {
+            for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable) {
+                for (ShardRouting shardRouting : indexShardRoutingTable) {
+                    counts.adjustOrPutValue(clusterState.nodes().get(shardRouting.currentNodeId()).name(), 1, 1);
+                }
+            }
+        }
+        assertThat(counts.get("node3"), equalTo(10));
+    }
+}