wip

crate · Mar 31, 2015 · b0b0ddf · b0b0ddf
1 parent 16c2473
commit b0b0ddf
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 32 deletions.
diff --git a/sql/src/test/java/io/crate/benchmark/LuceneDocCollectorBenchmark.java b/sql/src/test/java/io/crate/benchmark/LuceneDocCollectorBenchmark.java
@@ -57,6 +57,7 @@
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.index.service.IndexService;
 import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.indices.IndexMissingException;
 import org.elasticsearch.indices.IndicesService;
 import org.junit.Before;
 import org.junit.Rule;
@@ -124,42 +125,23 @@ public byte[] generateRowSource() throws IOException {
     }
 
     @Override
-    public Client getClient(boolean firstNode) {
-        return cluster.client(NODE1);
+    public boolean generateData() {
+        return true;
     }
 
     @Override
     @Before
     public void setUp() throws Exception {
-        if (NODE1 == null) {
-            NODE1 = cluster.startNode(getNodeSettings(1));
+        super.setUp();
+
+        IndexService indexService;
+        try {
+            IndicesService instanceFromNode = cluster.getInstanceFromNode(NODE2, IndicesService.class);
+            indexService = instanceFromNode.indexServiceSafe(INDEX_NAME);
+        } catch (IndexMissingException e) {
+            IndicesService instanceFromNode = cluster.getInstanceFromNode(NODE1, IndicesService.class);
+            indexService = instanceFromNode.indexServiceSafe(INDEX_NAME);
         }
-        if (!indexExists()) {
-            execute("create table \"" + INDEX_NAME + "\" (" +
-                    " \"areaInSqKm\" float," +
-                    " capital string," +
-                    " continent string," +
-                    " \"continentName\" string," +
-                    " \"countryCode\" string," +
-                    " \"countryName\" string," +
-                    " north float," +
-                    " east float," +
-                    " south float," +
-                    " west float," +
-                    " \"fipsCode\" string," +
-                    " \"currencyCode\" string," +
-                    " languages string," +
-                    " \"isoAlpha3\" string," +
-                    " \"isoNumeric\" string," +
-                    " population integer" +
-                    ") clustered into 1 shards with (number_of_replicas=0)", new Object[0], true);
-            client().admin().cluster().prepareHealth(INDEX_NAME).setWaitForGreenStatus().execute().actionGet();
-            refresh(client());
-        }
-        doGenerateData();
-
-        IndicesService instanceFromNode = cluster.getInstanceFromFirstNode(IndicesService.class);
-        IndexService indexService = instanceFromNode.indexServiceSafe(INDEX_NAME);
 
         shardCollectService = indexService.shardInjector(0).getInstance(ShardCollectService.class);
         collectContextService = indexService.shardInjector(0).getInstance(CollectContextService.class);
@@ -169,6 +151,29 @@ public void setUp() throws Exception {
         orderBy = new OrderBy(ImmutableList.of((Symbol) reference), new boolean[]{false}, new Boolean[]{false});
     }
 
+    @Override
+    protected void createTable() {
+        execute("create table \"" + INDEX_NAME + "\" (" +
+                " \"areaInSqKm\" float," +
+                " capital string," +
+                " continent string," +
+                " \"continentName\" string," +
+                " \"countryCode\" string," +
+                " \"countryName\" string," +
+                " north float," +
+                " east float," +
+                " south float," +
+                " west float," +
+                " \"fipsCode\" string," +
+                " \"currencyCode\" string," +
+                " languages string," +
+                " \"isoAlpha3\" string," +
+                " \"isoNumeric\" string," +
+                " population integer" +
+                ") clustered into 1 shards with (number_of_replicas=0)", new Object[0], true);
+        client().admin().cluster().prepareHealth(INDEX_NAME).setWaitForGreenStatus().execute().actionGet();
+    }
+
     private LuceneDocCollector createDocCollector(OrderBy orderBy, Integer limit, List<Symbol> input) throws Exception{
         return createDocCollector(orderBy, limit, collectingProjector, input);
     }

diff --git a/testing/src/main/java/io/crate/test/integration/CrateTestCluster.java b/testing/src/main/java/io/crate/test/integration/CrateTestCluster.java
@@ -112,11 +112,16 @@ public CrateTestCluster(long clusterSeed, String clusterName) {
         this(clusterSeed, 2, "network", clusterName, NodeSettingsSource.EMPTY);
     }
 
-    public <T> T getInstanceFromFirstNode(Class<T> clazz) {
-        NodeAndClient node = nodes.values().iterator().next();
+    public <T> T getInstanceFromNode(String nodeName, Class<T> clazz) {
+        NodeAndClient node = nodes.get(nodeName);
         return node.node.injector().getInstance(clazz);
     }
 
+    public <T> T getInstanceFromFirstNode(Class<T> clazz) {
+        String name = nodes.keySet().iterator().next();
+        return getInstanceFromNode(name, clazz);
+    }
+
     public CrateTestCluster(long clusterSeed, int numNodes, String mode, String clusterName, NodeSettingsSource nodeSettingsSource) {
         this.clusterName = clusterName;
         random = new Random(clusterSeed);