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.
This commit is contained in:
Alexander Reelsen 2013-12-09 14:39:22 +01:00
parent 7cd189b3d3
commit d2ee22f01c
1 changed files with 36 additions and 1 deletions

View File

@ -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 {