Several internal improvements to internal test cluster infra

core/src/main/java/org/elasticsearch/client/transport/TransportClient.java

@@ -235,7 +235,7 @@ ThreadPool getThreadPool() {
     public static final String CLIENT_TYPE = "transport";
 
     final Injector injector;
-    final NamedWriteableRegistry namedWriteableRegistry;
+    protected final NamedWriteableRegistry namedWriteableRegistry;
 
     private final List<LifecycleComponent> pluginLifecycleComponents;
     private final TransportClientNodesService nodesService;

test/framework/src/main/java/org/elasticsearch/common/settings/MockSecureSettings.java

@@ -38,6 +38,15 @@ public class MockSecureSettings implements SecureSettings {
     private Set<String> settingNames = new HashSet<>();
     private final AtomicBoolean closed = new AtomicBoolean(false);
 
+    public MockSecureSettings() {
+    }
+
+    private MockSecureSettings(MockSecureSettings source) {
+        secureStrings.putAll(source.secureStrings);
+        files.putAll(source.files);
+        settingNames.addAll(source.settingNames);
+    }
+
     @Override
     public boolean isLoaded() {
         return true;
@@ -94,4 +103,9 @@ private void ensureOpen() {
             throw new IllegalStateException("secure settings are already closed");
         }
     }
+
+    public SecureSettings clone() {
+        ensureOpen();
+        return new MockSecureSettings(this);
+    }
 }

test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java

@@ -1084,14 +1084,7 @@ protected void ensureClusterSizeConsistency() {
      */
     protected void ensureClusterStateConsistency() throws IOException {
         if (cluster() != null && cluster().size() > 0) {
-            final NamedWriteableRegistry namedWriteableRegistry;
-            if (isInternalCluster()) {
-                // If it's internal cluster - using existing registry in case plugin registered custom data
-                namedWriteableRegistry = internalCluster().getInstance(NamedWriteableRegistry.class);
-            } else {
-                // If it's external cluster - fall back to the standard set
-                namedWriteableRegistry = new NamedWriteableRegistry(ClusterModule.getNamedWriteables());
-            }
+            final NamedWriteableRegistry namedWriteableRegistry = cluster().getNamedWriteableRegistry();
             ClusterState masterClusterState = client().admin().cluster().prepareState().all().get().getState();
             byte[] masterClusterStateBytes = ClusterState.Builder.toBytes(masterClusterState);
             // remove local node reference
@@ -2192,9 +2185,13 @@ protected static RestClient createRestClient(RestClientBuilder.HttpClientConfigC
     }
 
     protected static RestClient createRestClient(RestClientBuilder.HttpClientConfigCallback httpClientConfigCallback, String protocol) {
-        final NodesInfoResponse nodeInfos = client().admin().cluster().prepareNodesInfo().get();
-        final List<NodeInfo> nodes = nodeInfos.getNodes();
-        assertFalse(nodeInfos.hasFailures());
+        NodesInfoResponse nodesInfoResponse = client().admin().cluster().prepareNodesInfo().get();
+        assertFalse(nodesInfoResponse.hasFailures());
+        return createRestClient(nodesInfoResponse.getNodes(), httpClientConfigCallback, protocol);
+    }
+
+    protected static RestClient createRestClient(final List<NodeInfo> nodes,
+                                                 RestClientBuilder.HttpClientConfigCallback httpClientConfigCallback, String protocol) {
         List<HttpHost> hosts = new ArrayList<>();
         for (NodeInfo node : nodes) {
             if (node.getHttp() != null) {

test/framework/src/main/java/org/elasticsearch/test/ExternalTestCluster.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.client.transport.TransportClient;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.common.network.NetworkModule;
 import org.elasticsearch.common.settings.Settings;
@@ -62,7 +63,7 @@ public final class ExternalTestCluster extends TestCluster {
     private static final AtomicInteger counter = new AtomicInteger();
     public static final String EXTERNAL_CLUSTER_PREFIX = "external_";
 
-    private final Client client;
+    private final MockTransportClient client;
 
     private final InetSocketAddress[] httpAddresses;
 
@@ -95,8 +96,7 @@ public ExternalTestCluster(Path tempDir, Settings additionalSettings, Collection
             }
         }
         Settings clientSettings = clientSettingsBuilder.build();
-        TransportClient client = new MockTransportClient(clientSettings, pluginClasses);
-
+        MockTransportClient client = new MockTransportClient(clientSettings, pluginClasses);
         try {
             client.addTransportAddresses(transportAddresses);
             NodesInfoResponse nodeInfos = client.admin().cluster().prepareNodesInfo().clear().setSettings(true).setHttp(true).get();
@@ -117,6 +117,7 @@ public ExternalTestCluster(Path tempDir, Settings additionalSettings, Collection
             this.numDataNodes = dataNodes;
             this.numMasterAndDataNodes = masterAndDataNodes;
             this.client = client;
+
             logger.info("Setup ExternalTestCluster [{}] made of [{}] nodes", nodeInfos.getClusterName().value(), size());
         } catch (Exception e) {
             client.close();
@@ -186,6 +187,11 @@ public Iterable<Client> getClients() {
         return Collections.singleton(client);
     }
 
+    @Override
+    public NamedWriteableRegistry getNamedWriteableRegistry() {
+        return client.getNamedWriteableRegistry();
+    }
+
     @Override
     public String getClusterName() {
         return clusterName;

test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java

@@ -51,9 +51,11 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.io.FileSystemUtils;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.common.network.NetworkModule;
+import org.elasticsearch.common.settings.MockSecureSettings;
 import org.elasticsearch.common.settings.SecureSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.Settings.Builder;
@@ -610,6 +612,10 @@ private NodeAndClient buildNode(int nodeId, long seed, Settings settings,
             throw new IllegalArgumentException(DISCOVERY_ZEN_MINIMUM_MASTER_NODES_SETTING.getKey() + " must be configured");
         }
         SecureSettings secureSettings = finalSettings.getSecureSettings();
+        if (secureSettings instanceof MockSecureSettings) {
+            // we clone this here since in the case of a node restart we might need it again
+            secureSettings = ((MockSecureSettings) secureSettings).clone();
+        }
         MockNode node = new MockNode(finalSettings.build(), plugins, nodeConfigurationSource.nodeConfigPath(nodeId));
         try {
             IOUtils.close(secureSettings);
@@ -1964,6 +1970,11 @@ public void remove() {
         };
     }
 
+    @Override
+    public NamedWriteableRegistry getNamedWriteableRegistry() {
+        return getInstance(NamedWriteableRegistry.class);
+    }
+
     /**
      * Returns a predicate that only accepts settings of nodes with one of the given names.
      */

test/framework/src/main/java/org/elasticsearch/test/TestCluster.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.metadata.IndexMetaData;
 import org.elasticsearch.cluster.metadata.IndexTemplateMetaData;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.indices.IndexTemplateMissingException;
@@ -34,6 +35,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.net.InetSocketAddress;
+import java.util.List;
 import java.util.Random;
 import java.util.Set;
 
@@ -233,5 +235,9 @@ public void wipeRepositories(String... repositories) {
      */
     public abstract Iterable<Client> getClients();
 
-
+    /**
+     * Returns this clusters {@link NamedWriteableRegistry} this is needed to
+     * deserialize binary content from this cluster that might include custom named writeables
+     */
+    public abstract NamedWriteableRegistry getNamedWriteableRegistry();
 }

test/framework/src/main/java/org/elasticsearch/transport/MockTransportClient.java

@@ -19,6 +19,7 @@
 package org.elasticsearch.transport;
 
 import org.elasticsearch.client.transport.TransportClient;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.network.NetworkModule;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.plugins.Plugin;
@@ -69,4 +70,8 @@ private static Collection<Class<? extends Plugin>> addMockTransportIfMissing(Set
         }
         return plugins;
     }
+
+    public NamedWriteableRegistry getNamedWriteableRegistry() {
+        return namedWriteableRegistry;
+    }
 }