[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.

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;

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,17 +57,17 @@
  * {@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 CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer;
 
-    public ClientYamlTestClient(
+    ClientYamlTestClient(
             final ClientYamlSuiteRestSpec restSpec,
             final RestClient restClient,
             final List<HttpHost> hosts,
@@ -195,10 +196,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();
         });
     }
@@ -240,4 +241,11 @@ private ClientYamlSuiteRestApi restApi(String apiName) {
         }
         return restApi;
     }
+
+    @Override
+    public void close() throws IOException {
+        for (RestClient restClient : restClients.values()) {
+            restClient.close();
+        }
+    }
 }

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

@@ -447,6 +447,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