From 30dca3178c3938d0414b88ec93e1b5076495833a Mon Sep 17 00:00:00 2001
From: javanna <cavannaluca@gmail.com>
Date: Tue, 26 Jun 2018 11:58:17 +0200
Subject: [PATCH] [TEST] Close additional clients created while running yaml
 tests

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.
---
 .../client/HasAttributeNodeSelector.java      | 19 +++++++++++++++++++
 .../test/rest/yaml/ClientYamlTestClient.java  | 18 +++++++++++++-----
 .../test/rest/yaml/section/DoSection.java     | 18 ++++++++++++++++++
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.java b/client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.java
index e4bb43458648b..11232a08c3d29 100644
--- a/client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.java
+++ b/client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.java
@@ -22,6 +22,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * A {@link NodeSelector} that selects nodes that have a particular value
@@ -49,6 +50,24 @@ public void select(Iterable<Node> nodes) {
         }
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        HasAttributeNodeSelector that = (HasAttributeNodeSelector) o;
+        return Objects.equals(key, that.key) &&
+                Objects.equals(value, that.value);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(key, value);
+    }
+
     @Override
     public String toString() {
         return key + "=" + value;
diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestClient.java b/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestClient.java
index fdc10a1a246e7..2d6bcc8cf5665 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestClient.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/ClientYamlTestClient.java
@@ -40,6 +40,7 @@
 import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestPath;
 import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestSpec;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.net.URI;
@@ -56,18 +57,18 @@
  * {@link RestClient} instance used to send the REST requests. Holds the {@link ClientYamlSuiteRestSpec} used to translate api calls into
  * REST calls.
  */
-public class ClientYamlTestClient {
+public class ClientYamlTestClient implements Closeable {
     private static final Logger logger = Loggers.getLogger(ClientYamlTestClient.class);
 
     private static final ContentType YAML_CONTENT_TYPE = ContentType.create("application/yaml");
 
     private final ClientYamlSuiteRestSpec restSpec;
-    protected final Map<NodeSelector, RestClient> restClients = new HashMap<>();
+    private final Map<NodeSelector, RestClient> restClients = new HashMap<>();
     private final Version esVersion;
     private final Version masterVersion;
     private final CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer;
 
-    public ClientYamlTestClient(
+    ClientYamlTestClient(
             final ClientYamlSuiteRestSpec restSpec,
             final RestClient restClient,
             final List<HttpHost> hosts,
@@ -202,10 +203,10 @@ protected RestClient getRestClient(NodeSelector nodeSelector) {
             RestClientBuilder builder = RestClient.builder(anyClient.getNodes().toArray(new Node[0]));
             try {
                 clientBuilderConsumer.accept(builder);
-            } catch(IOException e) {
+            } catch (IOException e) {
                 throw new UncheckedIOException(e);
             }
-            builder.setNodeSelector(nodeSelector);
+            builder.setNodeSelector(selector);
             return builder.build();
         });
     }
@@ -247,4 +248,11 @@ private ClientYamlSuiteRestApi restApi(String apiName) {
         }
         return restApi;
     }
+
+    @Override
+    public void close() throws IOException {
+        for (RestClient restClient : restClients.values()) {
+            restClient.close();
+        }
+    }
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java b/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java
index 8697b0bedcdf5..4e46a9ec89fd1 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/rest/yaml/section/DoSection.java
@@ -449,6 +449,24 @@ public void select(Iterable<Node> nodes) {
             lhs.select(nodes);
         }
 
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            ComposeNodeSelector that = (ComposeNodeSelector) o;
+            return Objects.equals(lhs, that.lhs) &&
+                    Objects.equals(rhs, that.rhs);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(lhs, rhs);
+        }
+
         @Override
         public String toString() {
             // . as in haskell's "compose" operator