Stableize SearchStatsTest after search refactoring

SearchStatsTest depends on a given set of nodes and shards. The test
needed to be adjusted to reflect a possibly random number of nodes.

src/test/java/org/elasticsearch/test/integration/AbstractSharedClusterTest.java

@@ -32,7 +32,9 @@ import org.elasticsearch.action.admin.indices.create.CreateIndexRequestBuilder;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.client.Requests;
 import org.elasticsearch.cluster.ClusterService;
+import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.MetaData;
+import org.elasticsearch.cluster.routing.GroupShardsIterator;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.settings.ImmutableSettings;
 import org.elasticsearch.common.settings.Settings;
@@ -252,4 +254,11 @@ public abstract class AbstractSharedClusterTest extends ElasticsearchTestCase {
     protected void optimize() {
         client().admin().indices().prepareOptimize().execute().actionGet();
     }
+    
+    protected int numAssignedShards(String... indices) {
+        ClusterState state = client().admin().cluster().prepareState().execute().actionGet().getState();
+        GroupShardsIterator allAssignedShardsGrouped = state.routingTable().allAssignedShardsGrouped(indices, true);
+        return allAssignedShardsGrouped.size();
+    }
+
 }

src/test/java/org/elasticsearch/test/integration/search/stats/SearchStatsTests.java

@@ -19,6 +19,13 @@
 
 package org.elasticsearch.test.integration.search.stats;
 
+import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
+
 import org.elasticsearch.action.admin.cluster.node.stats.NodesStatsResponse;
 import org.elasticsearch.action.admin.indices.stats.IndicesStatsResponse;
 import org.elasticsearch.action.search.SearchResponse;
@@ -27,13 +34,8 @@ import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.test.integration.AbstractSharedClusterTest;
-import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
-import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.*;
-
 /**
  */
 public class SearchStatsTests extends AbstractSharedClusterTest {
@@ -41,7 +43,6 @@ public class SearchStatsTests extends AbstractSharedClusterTest {
     @Override
     public Settings getSettings() {
         return randomSettingsBuilder()
-                .put("index.number_of_shards", 3)
                 .put("index.number_of_replicas", 0)
                 .build();
     }
@@ -52,16 +53,17 @@ public class SearchStatsTests extends AbstractSharedClusterTest {
         for (int i = 0; i < 500; i++) {
             client().prepareIndex("test1", "type", Integer.toString(i)).setSource("field", "value").execute().actionGet();
             if (i == 10) {
-                client().admin().indices().prepareRefresh("test1").execute().actionGet();
+                refresh();
             }
         }
         createIndex("test2");
         for (int i = 0; i < 500; i++) {
             client().prepareIndex("test2", "type", Integer.toString(i)).setSource("field", "value").execute().actionGet();
             if (i == 10) {
-                client().admin().indices().prepareRefresh("test1").execute().actionGet();
+                refresh();
             }
         }
+        cluster().ensureAtMostNumNodes(numAssignedShards("test1", "test2"));
         for (int i = 0; i < 200; i++) {
             client().prepareSearch().setQuery(QueryBuilders.termQuery("field", "value")).setStats("group1", "group2").execute().actionGet();
         }
@@ -102,7 +104,7 @@ public class SearchStatsTests extends AbstractSharedClusterTest {
                 .execute().actionGet();
 
         indicesStats = client().admin().indices().prepareStats().execute().actionGet();
-        assertThat(indicesStats.getTotal().getSearch().getOpenContexts(), equalTo(3l)); // 3 shards
+        assertThat(indicesStats.getTotal().getSearch().getOpenContexts(), equalTo((long)numAssignedShards("test1")));
 
         // scroll, but with no timeout (so no context)
         searchResponse = client().prepareSearchScroll(searchResponse.getScrollId()).execute().actionGet();