[Test] Do not rely on MockZenPing for Azure tests (#27945)

This commit changes some Azure tests so that they do not rely on
MockZenPing and TestZenDiscovery anymore, but instead use a mocked
AzureComputeService that exposes internal test cluster nodes as if
they were real Azure nodes.

Related to #27859

Closes #27917, #11533

Author: tlrx

core/src/main/java/org/elasticsearch/plugins/DiscoveryPlugin.java

@@ -19,7 +19,6 @@
 
 package org.elasticsearch.plugins;
 
-import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
 import java.util.function.BiConsumer;

plugins/discovery-azure-classic/src/main/java/org/elasticsearch/discovery/azure/classic/AzureUnicastHostsProvider.java

@@ -62,6 +62,10 @@ public enum HostType {
             this.type = type ;
         }
 
+        public String getType() {
+            return type;
+        }
+
         public static HostType fromString(String type) {
             for (HostType hostType : values()) {
                 if (hostType.type.equalsIgnoreCase(type)) {
@@ -196,43 +200,7 @@ public List<DiscoveryNode> buildDynamicNodes() {
             // In other case, it should be the right deployment so we can add it to the list of instances
 
             for (RoleInstance instance : deployment.getRoleInstances()) {
-                String networkAddress = null;
-                // Let's detect if we want to use public or private IP
-                switch (hostType) {
-                    case PRIVATE_IP:
-                        InetAddress privateIp = instance.getIPAddress();
-
-                        if (privateIp != null) {
-                            if (privateIp.equals(ipAddress)) {
-                                logger.trace("adding ourselves {}", NetworkAddress.format(ipAddress));
-                            }
-                            networkAddress = InetAddresses.toUriString(privateIp);
-                        } else {
-                            logger.trace("no private ip provided. ignoring [{}]...", instance.getInstanceName());
-                        }
-                        break;
-                    case PUBLIC_IP:
-                        for (InstanceEndpoint endpoint : instance.getInstanceEndpoints()) {
-                            if (!publicEndpointName.equals(endpoint.getName())) {
-                                logger.trace("ignoring endpoint [{}] as different than [{}]",
-                                        endpoint.getName(), publicEndpointName);
-                                continue;
-                            }
-
-                            networkAddress = NetworkAddress.format(new InetSocketAddress(endpoint.getVirtualIPAddress(),
-                                    endpoint.getPort()));
-                        }
-
-                        if (networkAddress == null) {
-                            logger.trace("no public ip provided. ignoring [{}]...", instance.getInstanceName());
-                        }
-                        break;
-                    default:
-                        // This could never happen!
-                        logger.warn("undefined host_type [{}]. Please check your settings.", hostType);
-                        return cachedDiscoNodes;
-                }
-
+                final String networkAddress = resolveInstanceAddress(hostType, instance);
                 if (networkAddress == null) {
                     // We have a bad parameter here or not enough information from azure
                     logger.warn("no network address found. ignoring [{}]...", instance.getInstanceName());
@@ -257,4 +225,24 @@ public List<DiscoveryNode> buildDynamicNodes() {
 
         return cachedDiscoNodes;
     }
+
+    protected String resolveInstanceAddress(final HostType hostType, final RoleInstance instance) {
+        if (hostType == HostType.PRIVATE_IP) {
+            final InetAddress privateIp = instance.getIPAddress();
+            if (privateIp != null) {
+                return InetAddresses.toUriString(privateIp);
+            } else {
+                logger.trace("no private ip provided. ignoring [{}]...", instance.getInstanceName());
+            }
+        } else if (hostType == HostType.PUBLIC_IP) {
+            for (InstanceEndpoint endpoint : instance.getInstanceEndpoints()) {
+                if (publicEndpointName.equals(endpoint.getName())) {
+                    return NetworkAddress.format(new InetSocketAddress(endpoint.getVirtualIPAddress(), endpoint.getPort()));
+                } else {
+                    logger.trace("ignoring endpoint [{}] as different than [{}]", endpoint.getName(), publicEndpointName);
+                }
+            }
+        }
+        return null;
+    }
 }

plugins/discovery-azure-classic/src/main/java/org/elasticsearch/plugin/discovery/azure/classic/AzureDiscoveryPlugin.java

@@ -61,10 +61,16 @@ protected AzureComputeService createComputeService() {
     public Map<String, Supplier<UnicastHostsProvider>> getZenHostsProviders(TransportService transportService,
                                                                             NetworkService networkService) {
         return Collections.singletonMap(AZURE,
-            () -> new AzureUnicastHostsProvider(settings, createComputeService(), transportService, networkService));
+            () -> createUnicastHostsProvider(settings, createComputeService(), transportService, networkService));
     }
 
-
+    // Used for testing
+    protected AzureUnicastHostsProvider createUnicastHostsProvider(final Settings settings,
+                                                                   final AzureComputeService azureComputeService,
+                                                                   final TransportService transportService,
+                                                                   final NetworkService networkService) {
+        return new AzureUnicastHostsProvider(settings, azureComputeService, transportService, networkService);
+    }
 
     @Override
     public List<Setting<?>> getSettings() {
@@ -79,6 +85,4 @@ public List<Setting<?>> getSettings() {
                             AzureComputeService.Discovery.DEPLOYMENT_SLOT_SETTING,
                             AzureComputeService.Discovery.ENDPOINT_NAME_SETTING);
     }
-
-
 }

plugins/discovery-azure-classic/src/test/java/org/elasticsearch/cloud/azure/classic/AbstractAzureComputeServiceTestCase.java

@@ -19,28 +19,45 @@
 
 package org.elasticsearch.cloud.azure.classic;
 
+import com.microsoft.windowsazure.management.compute.models.DeploymentSlot;
+import com.microsoft.windowsazure.management.compute.models.DeploymentStatus;
+import com.microsoft.windowsazure.management.compute.models.HostedServiceGetDetailedResponse;
+import com.microsoft.windowsazure.management.compute.models.InstanceEndpoint;
+import com.microsoft.windowsazure.management.compute.models.RoleInstance;
+import com.microsoft.windowsazure.management.compute.models.RoleInstancePowerState;
 import org.elasticsearch.action.admin.cluster.node.info.NodesInfoResponse;
+import org.elasticsearch.cloud.azure.classic.management.AzureComputeService;
 import org.elasticsearch.cloud.azure.classic.management.AzureComputeService.Discovery;
 import org.elasticsearch.cloud.azure.classic.management.AzureComputeService.Management;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.network.NetworkAddress;
+import org.elasticsearch.common.network.NetworkService;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.transport.TransportAddress;
+import org.elasticsearch.discovery.azure.classic.AzureUnicastHostsProvider;
+import org.elasticsearch.discovery.zen.ZenDiscovery;
+import org.elasticsearch.plugin.discovery.azure.classic.AzureDiscoveryPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESIntegTestCase;
+import org.elasticsearch.transport.TransportService;
+import org.junit.After;
 
-import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static org.elasticsearch.common.util.CollectionUtils.newSingletonArrayList;
 
 public abstract class AbstractAzureComputeServiceTestCase extends ESIntegTestCase {
 
-    private Class<? extends Plugin> mockPlugin;
+    private static final Map<String, DiscoveryNode> nodes = new ConcurrentHashMap<>();
 
-    public AbstractAzureComputeServiceTestCase(Class<? extends Plugin> mockPlugin) {
-        // We want to inject the Azure API Mock
-        this.mockPlugin = mockPlugin;
-    }
-
-    @Override
-    protected boolean addTestZenDiscovery() {
-        return false;
+    @After
+    public void clearAzureNodes() {
+        nodes.clear();
     }
 
     @Override
@@ -49,6 +66,10 @@ protected Settings nodeSettings(int nodeOrdinal) {
             .put(super.nodeSettings(nodeOrdinal))
             .put("discovery.zen.hosts_provider", "azure");
 
+        // Make the test run faster
+        builder.put(ZenDiscovery.JOIN_TIMEOUT_SETTING.getKey(), "1s")
+            .put(ZenDiscovery.PING_TIMEOUT_SETTING.getKey(), "500ms");
+
         // We add a fake subscription_id to start mock compute service
         builder.put(Management.SUBSCRIPTION_ID_SETTING.getKey(), "fake")
             .put(Discovery.REFRESH_SETTING.getKey(), "5s")
@@ -60,13 +81,107 @@ protected Settings nodeSettings(int nodeOrdinal) {
 
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
-        return Arrays.asList(mockPlugin);
+        return Collections.singletonList(TestPlugin.class);
     }
 
-    protected void checkNumberOfNodes(int expected) {
-        NodesInfoResponse nodeInfos = client().admin().cluster().prepareNodesInfo().execute().actionGet();
+    @Override
+    protected boolean addTestZenDiscovery() {
+        return false;
+    }
+
+    /**
+     * Register an existing node as a Azure node, exposing its address and details htrough
+     *
+     * @param nodeName the name of the node
+     */
+    protected void registerAzureNode(final String nodeName) {
+        TransportService transportService = internalCluster().getInstance(TransportService.class, nodeName);
+        assertNotNull(transportService);
+        DiscoveryNode discoveryNode = transportService.getLocalNode();
+        assertNotNull(discoveryNode);
+        if (nodes.put(discoveryNode.getName(), discoveryNode) != null) {
+            throw new IllegalArgumentException("Node [" + discoveryNode.getName() + "] cannot be registered twice in Azure");
+        }
+    }
+
+    protected void assertNumberOfNodes(int expected) {
+        NodesInfoResponse nodeInfos = client().admin().cluster().prepareNodesInfo().clear().execute().actionGet();
         assertNotNull(nodeInfos);
         assertNotNull(nodeInfos.getNodes());
         assertEquals(expected, nodeInfos.getNodes().size());
     }
+
+    /**
+     * Test plugin that exposes internal test cluster nodes as if they were real Azure nodes.
+     * Use {@link #registerAzureNode(String)} method to expose nodes in the tests.
+     */
+    public static class TestPlugin extends AzureDiscoveryPlugin {
+
+        public TestPlugin(Settings settings) {
+            super(settings);
+        }
+
+        @Override
+        protected AzureComputeService createComputeService() {
+            return () -> {
+                final List<RoleInstance> instances = new ArrayList<>();
+                for (Map.Entry<String, DiscoveryNode> node : nodes.entrySet()) {
+                    final String name = node.getKey();
+                    final DiscoveryNode discoveryNode = node.getValue();
+
+                    RoleInstance instance = new RoleInstance();
+                    instance.setInstanceName(name);
+                    instance.setHostName(discoveryNode.getHostName());
+                    instance.setPowerState(RoleInstancePowerState.Started);
+
+                    // Set the private IP address
+                    final TransportAddress transportAddress = discoveryNode.getAddress();
+                    instance.setIPAddress(transportAddress.address().getAddress());
+
+                    // Set the public IP address
+                    final InstanceEndpoint endpoint = new InstanceEndpoint();
+                    endpoint.setName(Discovery.ENDPOINT_NAME_SETTING.getDefault(Settings.EMPTY));
+                    endpoint.setVirtualIPAddress(transportAddress.address().getAddress());
+                    endpoint.setPort(transportAddress.address().getPort());
+                    instance.setInstanceEndpoints(new ArrayList<>(Collections.singletonList(endpoint)));
+                    instances.add(instance);
+                }
+
+                final HostedServiceGetDetailedResponse.Deployment deployment = new HostedServiceGetDetailedResponse.Deployment();
+                deployment.setName("dummy");
+                deployment.setDeploymentSlot(DeploymentSlot.Production);
+                deployment.setStatus(DeploymentStatus.Running);
+                deployment.setRoleInstances(new ArrayList<>(Collections.unmodifiableList(instances)));
+
+                final HostedServiceGetDetailedResponse response = new HostedServiceGetDetailedResponse();
+                response.setDeployments(newSingletonArrayList(deployment));
+
+                return response;
+            };
+        }
+
+        /**
+         * Defines a {@link AzureUnicastHostsProvider} for testing purpose that is able to resolve
+         * network addresses for Azure instances running on the same host but different ports.
+         */
+        @Override
+        protected AzureUnicastHostsProvider createUnicastHostsProvider(final Settings settings,
+                                                                       final AzureComputeService azureComputeService,
+                                                                       final TransportService transportService,
+                                                                       final NetworkService networkService) {
+            return new AzureUnicastHostsProvider(settings, azureComputeService, transportService, networkService) {
+                @Override
+                protected String resolveInstanceAddress(final HostType hostType, final RoleInstance instance) {
+                    if (hostType == HostType.PRIVATE_IP) {
+                        DiscoveryNode discoveryNode = nodes.get(instance.getInstanceName());
+                        if (discoveryNode != null) {
+                            // Format the InetSocketAddress to a format that contains the port number
+                            return NetworkAddress.format(discoveryNode.getAddress().address());
+                        }
+                    }
+                    return super.resolveInstanceAddress(hostType, instance);
+                }
+            };
+        }
+    }
 }