Add unit tests for ec2 AZ node attributes

plugins/discovery-ec2/src/main/java/org/elasticsearch/plugin/discovery/ec2/Ec2DiscoveryPlugin.java

@@ -143,6 +143,11 @@ public class Ec2DiscoveryPlugin extends Plugin {
     /** Adds a node attribute for the ec2 availability zone. */
     @Override
     public Settings additionalSettings() {
+        return getAvailabilityZoneNodeAttributes(settings, AwsEc2ServiceImpl.EC2_METADATA_URL + "placement/availability-zone");
+    }
+
+    // pkg private for testing
+    static Settings getAvailabilityZoneNodeAttributes(Settings settings, String azMetadataUrl) {
         if (AwsEc2Service.AUTO_ATTRIBUTE_SETTING.get(settings) == false) {
             return Settings.EMPTY;
         }
@@ -151,7 +156,7 @@ public class Ec2DiscoveryPlugin extends Plugin {
         final URL url;
         final URLConnection urlConnection;
         try {
-            url = new URL(AwsEc2ServiceImpl.EC2_METADATA_URL + "placement/availability-zone");
+            url = new URL(azMetadataUrl);
             logger.debug("obtaining ec2 [placement/availability-zone] from ec2 meta-data url {}", url);
             urlConnection = url.openConnection();
             urlConnection.setConnectTimeout(2000);
@@ -165,7 +170,7 @@ public class Ec2DiscoveryPlugin extends Plugin {
 
             String metadataResult = urlReader.readLine();
             if (metadataResult == null || metadataResult.length() == 0) {
-                throw new IOException("no ec2 metadata returned from " + url);
+                throw new IllegalStateException("no ec2 metadata returned from " + url);
             } else {
                 attrs.put(Node.NODE_ATTRIBUTES.getKey() + "aws_availability_zone", metadataResult);
             }

plugins/discovery-ec2/src/test/java/org/elasticsearch/plugin/discovery/ec2/Ec2DiscoveryPluginTests.java

@@ -0,0 +1,83 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.plugin.discovery.ec2;
+
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Arrays;
+
+import org.elasticsearch.cloud.aws.AwsEc2Service;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.node.Node;
+import org.elasticsearch.test.ESTestCase;
+
+public class Ec2DiscoveryPluginTests extends ESTestCase {
+
+    private Settings getNodeAttributes(Settings settings, String url) {
+        Settings realSettings = Settings.builder()
+            .put(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), true)
+            .put(settings).build();
+        return Ec2DiscoveryPlugin.getAvailabilityZoneNodeAttributes(realSettings, url);
+    }
+
+    private void assertNodeAttributes(Settings settings, String url, String expected) {
+        Settings additional = getNodeAttributes(settings, url);
+        if (expected == null) {
+            assertTrue(additional.isEmpty());
+        } else {
+            assertEquals(expected, additional.get(Node.NODE_ATTRIBUTES.getKey() + "aws_availability_zone"));
+        }
+    }
+
+    public void testNodeAttributesDisabled() {
+        Settings settings = Settings.builder()
+            .put(AwsEc2Service.AUTO_ATTRIBUTE_SETTING.getKey(), false).build();
+        assertNodeAttributes(settings, "bogus", null);
+    }
+
+    public void testNodeAttributes() throws Exception {
+        Path zoneUrl = createTempFile();
+        Files.write(zoneUrl, Arrays.asList("us-east-1c"));
+        assertNodeAttributes(Settings.EMPTY, zoneUrl.toUri().toURL().toString(), "us-east-1c");
+    }
+
+    public void testNodeAttributesBogusUrl() {
+        UncheckedIOException e = expectThrows(UncheckedIOException.class, () ->
+            getNodeAttributes(Settings.EMPTY, "bogus")
+        );
+        assertNotNull(e.getCause());
+        String msg = e.getCause().getMessage();
+        assertTrue(msg, msg.contains("no protocol: bogus"));
+    }
+
+    public void testNodeAttributesEmpty() throws Exception {
+        Path zoneUrl = createTempFile();
+        IllegalStateException e = expectThrows(IllegalStateException.class, () ->
+            getNodeAttributes(Settings.EMPTY, zoneUrl.toUri().toURL().toString())
+        );
+        assertTrue(e.getMessage(), e.getMessage().contains("no ec2 metadata returned"));
+    }
+
+    public void testNodeAttributesErrorLenient() throws Exception {
+        Path dne = createTempDir().resolve("dne");
+        assertNodeAttributes(Settings.EMPTY, dne.toUri().toURL().toString(), null);
+    }
+}