Skip to content

Commit e6d3249

Browse files
committed
[TEST] Close additional clients created while running yaml tests (#31575)
We recently introduced a mechanism that allows to specify a node selector as part of do sections (see #31471). When a node selector that is not the default one is configured, a new client will be initialized with the same properties as the default one, but with the specified node selector. This commit improves such mechanism but also closing the additional clients being created and adding equals/hashcode impl to the custom node selector as they are cached into a map.
1 parent 675cbdc commit e6d3249

File tree

3 files changed

+50
-5
lines changed

3 files changed

+50
-5
lines changed

client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.java

+19
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import java.util.Iterator;
2323
import java.util.List;
2424
import java.util.Map;
25+
import java.util.Objects;
2526

2627
/**
2728
* A {@link NodeSelector} that selects nodes that have a particular value
@@ -49,6 +50,24 @@ public void select(Iterable<Node> nodes) {
4950
}
5051
}
5152

53+
@Override
54+
public boolean equals(Object o) {
55+
if (this == o) {
56+
return true;
57+
}
58+
if (o == null || getClass() != o.getClass()) {
59+
return false;
60+
}
61+
HasAttributeNodeSelector that = (HasAttributeNodeSelector) o;
62+
return Objects.equals(key, that.key) &&
63+
Objects.equals(value, that.value);
64+
}
65+
66+
@Override
67+
public int hashCode() {
68+
return Objects.hash(key, value);
69+
}
70+
5271
@Override
5372
public String toString() {
5473
return key + "=" + value;

test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestClient.java

+13-5
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestPath;
4141
import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestSpec;
4242

43+
import java.io.Closeable;
4344
import java.io.IOException;
4445
import java.io.UncheckedIOException;
4546
import java.net.URI;
@@ -56,17 +57,17 @@
5657
* {@link RestClient} instance used to send the REST requests. Holds the {@link ClientYamlSuiteRestSpec} used to translate api calls into
5758
* REST calls.
5859
*/
59-
public class ClientYamlTestClient {
60+
public class ClientYamlTestClient implements Closeable {
6061
private static final Logger logger = Loggers.getLogger(ClientYamlTestClient.class);
6162

6263
private static final ContentType YAML_CONTENT_TYPE = ContentType.create("application/yaml");
6364

6465
private final ClientYamlSuiteRestSpec restSpec;
65-
protected final Map<NodeSelector, RestClient> restClients = new HashMap<>();
66+
private final Map<NodeSelector, RestClient> restClients = new HashMap<>();
6667
private final Version esVersion;
6768
private final CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer;
6869

69-
public ClientYamlTestClient(
70+
ClientYamlTestClient(
7071
final ClientYamlSuiteRestSpec restSpec,
7172
final RestClient restClient,
7273
final List<HttpHost> hosts,
@@ -195,10 +196,10 @@ protected RestClient getRestClient(NodeSelector nodeSelector) {
195196
RestClientBuilder builder = RestClient.builder(anyClient.getNodes().toArray(new Node[0]));
196197
try {
197198
clientBuilderConsumer.accept(builder);
198-
} catch(IOException e) {
199+
} catch (IOException e) {
199200
throw new UncheckedIOException(e);
200201
}
201-
builder.setNodeSelector(nodeSelector);
202+
builder.setNodeSelector(selector);
202203
return builder.build();
203204
});
204205
}
@@ -240,4 +241,11 @@ private ClientYamlSuiteRestApi restApi(String apiName) {
240241
}
241242
return restApi;
242243
}
244+
245+
@Override
246+
public void close() throws IOException {
247+
for (RestClient restClient : restClients.values()) {
248+
restClient.close();
249+
}
250+
}
243251
}

test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java

+18
Original file line numberDiff line numberDiff line change
@@ -447,6 +447,24 @@ public void select(Iterable<Node> nodes) {
447447
lhs.select(nodes);
448448
}
449449

450+
@Override
451+
public boolean equals(Object o) {
452+
if (this == o) {
453+
return true;
454+
}
455+
if (o == null || getClass() != o.getClass()) {
456+
return false;
457+
}
458+
ComposeNodeSelector that = (ComposeNodeSelector) o;
459+
return Objects.equals(lhs, that.lhs) &&
460+
Objects.equals(rhs, that.rhs);
461+
}
462+
463+
@Override
464+
public int hashCode() {
465+
return Objects.hash(lhs, rhs);
466+
}
467+
450468
@Override
451469
public String toString() {
452470
// . as in haskell's "compose" operator

0 commit comments

Comments
 (0)