From d2ee22f01c929a9b62293af51f33e8e5072c1c08 Mon Sep 17 00:00:00 2001
From: Alexander Reelsen <alexander@reelsen.net>
Date: Mon, 9 Dec 2013 14:39:22 +0100
Subject: [PATCH] Allow to get a specific transport client inside of tests

Adding functionality to call cluster().transportClient() in tests in order
to get an arbitrary TransportClient object back, independently if the
transport client ratio in returning the normal clients is configured.

Also made sure, that if the normal client is already a transport client
(or a node client) we do not generate another one.
---
 .../org/elasticsearch/test/TestCluster.java   | 37 ++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/test/java/org/elasticsearch/test/TestCluster.java b/src/test/java/org/elasticsearch/test/TestCluster.java
index aa52129d7e4..101a18fbd49 100644
--- a/src/test/java/org/elasticsearch/test/TestCluster.java
+++ b/src/test/java/org/elasticsearch/test/TestCluster.java
@@ -26,6 +26,7 @@ import com.google.common.collect.Iterators;
 import com.google.common.collect.Sets;
 import org.apache.lucene.util.IOUtils;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.client.node.NodeClient;
 import org.elasticsearch.client.transport.TransportClient;
 import org.elasticsearch.cluster.ClusterService;
 import org.elasticsearch.cluster.ClusterState;
@@ -321,6 +322,15 @@ public final class TestCluster implements Iterable<Client> {
         return getRandomNodeAndClient(new ClientNodePredicate()).client(random);
     }
 
+    /**
+     * Returns a transport client
+     */
+    public synchronized Client transportClient() {
+        ensureOpen();
+        // randomly return a transport client going to one of the nodes in the cluster
+        return getOrBuildRandomNode().transportClient();
+    }
+
     /**
      * Returns a node client to a given node.
      */
@@ -378,6 +388,7 @@ public final class TestCluster implements Iterable<Client> {
         private InternalNode node;
         private Client client;
         private Client nodeClient;
+        private Client transportClient;
         private final AtomicBoolean closed = new AtomicBoolean(false);
         private final ClientFactory clientFactory;
         private final String name;
@@ -410,11 +421,31 @@ public final class TestCluster implements Iterable<Client> {
                 throw new RuntimeException("already closed");
             }
             if (nodeClient == null) {
-                nodeClient = node.client();
+                Client maybeNodeClient = client(random);
+                if (client instanceof NodeClient) {
+                    nodeClient = maybeNodeClient;
+                } else {
+                    nodeClient = node.client();
+                }
             }
             return nodeClient;
         }
 
+        Client transportClient() {
+            if (closed.get()) {
+                throw new RuntimeException("already closed");
+            }
+            if (transportClient == null) {
+                Client maybeTransportClient = client(random);
+                if (maybeTransportClient instanceof TransportClient) {
+                    transportClient = maybeTransportClient;
+                } else {
+                    transportClient = TransportClientFactory.NO_SNIFF_CLIENT_FACTORY.client(node, clusterName, random);
+                }
+            }
+            return transportClient;
+        }
+
         void resetClient() {
             if (closed.get()) {
                 throw new RuntimeException("already closed");
@@ -427,6 +458,10 @@ public final class TestCluster implements Iterable<Client> {
                 nodeClient.close();
                 nodeClient = null;
             }
+            if (transportClient != null) {
+                transportClient.close();
+                transportClient = null;
+            }
         }
 
         void restart(RestartCallback callback) throws Exception {