Skip to content

Commit 823a9d3

Browse files
authored
[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 8a6d062 commit 823a9d3

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

Lines changed: 19 additions & 0 deletions
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

Lines changed: 13 additions & 5 deletions
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,18 +57,18 @@
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 Version masterVersion;
6869
private final CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer;
6970

70-
public ClientYamlTestClient(
71+
ClientYamlTestClient(
7172
final ClientYamlSuiteRestSpec restSpec,
7273
final RestClient restClient,
7374
final List<HttpHost> hosts,
@@ -202,10 +203,10 @@ protected RestClient getRestClient(NodeSelector nodeSelector) {
202203
RestClientBuilder builder = RestClient.builder(anyClient.getNodes().toArray(new Node[0]));
203204
try {
204205
clientBuilderConsumer.accept(builder);
205-
} catch(IOException e) {
206+
} catch (IOException e) {
206207
throw new UncheckedIOException(e);
207208
}
208-
builder.setNodeSelector(nodeSelector);
209+
builder.setNodeSelector(selector);
209210
return builder.build();
210211
});
211212
}
@@ -247,4 +248,11 @@ private ClientYamlSuiteRestApi restApi(String apiName) {
247248
}
248249
return restApi;
249250
}
251+
252+
@Override
253+
public void close() throws IOException {
254+
for (RestClient restClient : restClients.values()) {
255+
restClient.close();
256+
}
257+
}
250258
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,24 @@ public void select(Iterable<Node> nodes) {
449449
lhs.select(nodes);
450450
}
451451

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

0 commit comments

Comments
 (0)