Test: Only sniff host metadata for node_selectors

docs/src/test/java/org/elasticsearch/smoketest/DocsClientYamlTestSuiteIT.java

@@ -92,8 +92,7 @@ protected ClientYamlTestClient initClientYamlTestClient(
             final RestClient restClient,
             final List<HttpHost> hosts,
             final Version esVersion) {
-        return new ClientYamlDocsTestClient(restSpec, restClient, hosts, esVersion,
-                restClientBuilder -> configureClient(restClientBuilder, restClientSettings()));
+        return new ClientYamlDocsTestClient(restSpec, restClient, hosts, esVersion, this::getClientBuilderWithSniffedHosts);
     }
 
     /**

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

@@ -28,7 +28,7 @@
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpHost;
 import org.elasticsearch.client.RestClientBuilder;
-import org.elasticsearch.common.CheckedConsumer;
+import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestSpec;
 
 import java.io.IOException;
@@ -49,8 +49,8 @@ public ClientYamlDocsTestClient(
             final RestClient restClient,
             final List<HttpHost> hosts,
             final Version esVersion,
-            final CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer) {
-        super(restSpec, restClient, hosts, esVersion, clientBuilderConsumer);
+            final CheckedSupplier<RestClientBuilder, IOException> clientBuilderWithSniffedNodes) {
+        super(restSpec, restClient, hosts, esVersion, clientBuilderWithSniffedNodes);
     }
 
     @Override

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

@@ -26,15 +26,14 @@
 import org.apache.http.util.EntityUtils;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.Version;
-import org.elasticsearch.client.Node;
 import org.elasticsearch.client.NodeSelector;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.client.RestClientBuilder;
-import org.elasticsearch.common.CheckedConsumer;
+import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestApi;
 import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestPath;
@@ -65,19 +64,19 @@ public class ClientYamlTestClient implements Closeable {
     private final ClientYamlSuiteRestSpec restSpec;
     private final Map<NodeSelector, RestClient> restClients = new HashMap<>();
     private final Version esVersion;
-    private final CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer;
+    private final CheckedSupplier<RestClientBuilder, IOException> clientBuilderWithSniffedNodes;
 
     ClientYamlTestClient(
             final ClientYamlSuiteRestSpec restSpec,
             final RestClient restClient,
             final List<HttpHost> hosts,
             final Version esVersion,
-            final CheckedConsumer<RestClientBuilder, IOException> clientBuilderConsumer) {
+            final CheckedSupplier<RestClientBuilder, IOException> clientBuilderWithSniffedNodes) {
         assert hosts.size() > 0;
         this.restSpec = restSpec;
         this.restClients.put(NodeSelector.ANY, restClient);
         this.esVersion = esVersion;
-        this.clientBuilderConsumer = clientBuilderConsumer;
+        this.clientBuilderWithSniffedNodes = clientBuilderWithSniffedNodes;
     }
 
     public Version getEsVersion() {
@@ -192,10 +191,9 @@ public ClientYamlTestResponse callApi(String apiName, Map<String, String> params
     protected RestClient getRestClient(NodeSelector nodeSelector) {
         //lazily build a new client in case we need to point to some specific node
         return restClients.computeIfAbsent(nodeSelector, selector -> {
-            RestClient anyClient = restClients.get(NodeSelector.ANY);
-            RestClientBuilder builder = RestClient.builder(anyClient.getNodes().toArray(new Node[0]));
+            RestClientBuilder builder;
             try {
-                clientBuilderConsumer.accept(builder);
+                builder = clientBuilderWithSniffedNodes.get();
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
             }

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

@@ -27,6 +27,7 @@
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.RestClientBuilder;
 import org.elasticsearch.client.sniff.ElasticsearchNodesSniffer;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Tuple;
@@ -58,13 +59,6 @@
 /**
  * Runs a suite of yaml tests shared with all the official Elasticsearch
  * clients against against an elasticsearch cluster.
- * <p>
- * <strong>IMPORTANT</strong>: These tests sniff the cluster for metadata
- * and hosts on startup and replace the list of hosts that they are
- * configured to use with the list sniffed from the cluster. So you can't
- * control which nodes receive the request by providing the right list of
- * nodes in the <code>tests.rest.cluster</code> system property. Instead
- * the tests must explictly use `node_selector`s.
  */
 public abstract class ESClientYamlSuiteTestCase extends ESRestTestCase {
 
@@ -116,11 +110,6 @@ protected ESClientYamlSuiteTestCase(ClientYamlTestCandidate testCandidate) {
     @Before
     public void initAndResetContext() throws Exception {
         if (restTestExecutionContext == null) {
-            // Sniff host metadata in case we need it in the yaml tests
-            List<Node> nodesWithMetadata = sniffHostMetadata();
-            client().setNodes(nodesWithMetadata);
-            adminClient().setNodes(nodesWithMetadata);
-
             assert adminExecutionContext == null;
             assert blacklistPathMatchers == null;
             ClientYamlSuiteRestSpec restSpec = ClientYamlSuiteRestSpec.load(SPEC_PATH);
@@ -172,8 +161,7 @@ protected ClientYamlTestClient initClientYamlTestClient(
             final RestClient restClient,
             final List<HttpHost> hosts,
             final Version esVersion) {
-        return new ClientYamlTestClient(restSpec, restClient, hosts, esVersion,
-                restClientBuilder -> configureClient(restClientBuilder, restClientSettings()));
+        return new ClientYamlTestClient(restSpec, restClient, hosts, esVersion, this::getClientBuilderWithSniffedHosts);
     }
 
     @AfterClass
@@ -407,13 +395,16 @@ protected boolean randomizeContentType() {
     }
 
     /**
-     * Sniff the cluster for host metadata.
+     * Sniff the cluster for host metadata and return a
+     * {@link RestClientBuilder} for a client with that metadata.
      */
-    private List<Node> sniffHostMetadata() throws IOException {
+    protected final RestClientBuilder getClientBuilderWithSniffedHosts() throws IOException {
         ElasticsearchNodesSniffer.Scheme scheme =
             ElasticsearchNodesSniffer.Scheme.valueOf(getProtocol().toUpperCase(Locale.ROOT));
         ElasticsearchNodesSniffer sniffer = new ElasticsearchNodesSniffer(
                 adminClient(), ElasticsearchNodesSniffer.DEFAULT_SNIFF_REQUEST_TIMEOUT, scheme);
-        return sniffer.sniff();
+        RestClientBuilder builder = RestClient.builder(sniffer.sniff().toArray(new Node[0]));
+        configureClient(builder, restClientSettings());
+        return builder;
     }
 }

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

@@ -389,7 +389,32 @@ private static NodeSelector parseAttributeValuesSelector(XContentParser parser)
             if (token == XContentParser.Token.FIELD_NAME) {
                 key = parser.currentName();
             } else if (token.isValue()) {
-                NodeSelector newSelector = new HasAttributeNodeSelector(key, parser.text());
+                /*
+                 * HasAttributeNodeSelector selects nodes that do not have
+                 * attribute metadata set so it can be used against nodes that
+                 * have not yet been sniffed. In these tests we expect the node
+                 * metadata to be explicitly sniffed if we need it and we'd
+                 * like to hard fail if it is not so we wrap the selector so we
+                 * can assert that the data is sniffed.
+                 */
+                NodeSelector delegate = new HasAttributeNodeSelector(key, parser.text());
+                NodeSelector newSelector = new NodeSelector() {
+                    @Override
+                    public void select(Iterable<Node> nodes) {
+                        for (Node node : nodes) {
+                            if (node.getAttributes() == null) {
+                                throw new IllegalStateException("expected [attributes] metadata to be set but got "
+                                        + node);
+                            }
+                        }
+                        delegate.select(nodes);
+                    }
+
+                    @Override
+                    public String toString() {
+                        return delegate.toString();
+                    }
+                };
                 result = result == NodeSelector.ANY ?
                     newSelector : new ComposeNodeSelector(result, newSelector);
             } else {

test/framework/src/test/java/org/elasticsearch/test/rest/yaml/section/DoSectionTests.java

@@ -539,6 +539,15 @@ public void testNodeSelectorByVersion() throws IOException {
         doSection.execute(context);
         verify(context).callApi("indices.get_field_mapping", singletonMap("index", "test_index"),
                 emptyList(), emptyMap(), doSection.getApiCallSection().getNodeSelector());
+
+        {
+            List<Node> badNodes = new ArrayList<>();
+            badNodes.add(new Node(new HttpHost("dummy")));
+            Exception e = expectThrows(IllegalStateException.class, () ->
+                    doSection.getApiCallSection().getNodeSelector().select(badNodes));
+            assertEquals("expected [version] metadata to be set but got [host=http://dummy]",
+                    e.getMessage());
+        }
     }
 
     private static Node nodeWithVersion(String version) {
@@ -567,6 +576,14 @@ public void testNodeSelectorByAttribute() throws IOException {
             doSection.getApiCallSection().getNodeSelector().select(nodes);
             assertEquals(Arrays.asList(hasAttr), nodes);
         }
+        {
+            List<Node> badNodes = new ArrayList<>();
+            badNodes.add(new Node(new HttpHost("dummy")));
+            Exception e = expectThrows(IllegalStateException.class, () ->
+                    doSection.getApiCallSection().getNodeSelector().select(badNodes));
+            assertEquals("expected [attributes] metadata to be set but got [host=http://dummy]",
+                    e.getMessage());
+        }
 
         parser = createParser(YamlXContent.yamlXContent,
                 "node_selector:\n" +

x-pack/docs/src/test/java/org/elasticsearch/smoketest/XDocsClientYamlTestSuiteIT.java

@@ -58,8 +58,7 @@ protected ClientYamlTestClient initClientYamlTestClient(
             final RestClient restClient,
             final List<HttpHost> hosts,
             final Version esVersion) {
-        return new ClientYamlDocsTestClient(restSpec, restClient, hosts, esVersion,
-                restClientBuilder -> configureClient(restClientBuilder, restClientSettings()));
+        return new ClientYamlDocsTestClient(restSpec, restClient, hosts, esVersion, this::getClientBuilderWithSniffedHosts);
     }
 
     /**