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:
parent
7cd189b3d3
commit
d2ee22f01c
|
@ -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 {
|
||||
|
|
Loading…
Reference in New Issue