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 e4bb4345864..11232a08c3d 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 @@ package org.elasticsearch.client;
 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 final class HasAttributeNodeSelector implements NodeSelector {
         }
     }
 
+    @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 fdc10a1a246..2d6bcc8cf56 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.ClientYamlSuiteRestApi;
 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 @@ import java.util.stream.Collectors;
  * {@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 @@ public class ClientYamlTestClient {
             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 @@ public class ClientYamlTestClient {
         }
         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 8697b0bedcd..4e46a9ec89f 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 class DoSection implements ExecutableSection {
             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