Add tests

Signed-off-by: Yanlong He <[email protected]>
opensearch-project · Dec 12, 2023 · 6a68a0f · 6a68a0f
1 parent bbbb93f
commit 6a68a0f
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 7 deletions.
diff --git a/plugins/discovery-ec2/src/main/java/org/opensearch/discovery/ec2/AwsEc2ServiceImpl.java b/plugins/discovery-ec2/src/main/java/org/opensearch/discovery/ec2/AwsEc2ServiceImpl.java
@@ -99,13 +99,7 @@ protected Ec2Client buildClient(
 
         if (Strings.hasText(endpoint)) {
             logger.debug("using explicit ec2 endpoint [{}]", endpoint);
-            URI uri = URI.create(endpoint);
-            if (uri.getScheme() == null) {
-                // if no scheme is provided, default to https
-                logger.debug("no scheme found in endpoint [{}], defaulting to https", endpoint);
-                uri = URI.create("https://" + endpoint);
-            }
-            builder.endpointOverride(uri);
+            builder.endpointOverride(URI.create(getFullEndpoint(endpoint)));
         }
 
         if (Strings.hasText(region)) {
@@ -116,6 +110,19 @@ protected Ec2Client buildClient(
         return SocketAccess.doPrivileged(builder::build);
     }
 
+    protected String getFullEndpoint(String endpoint) {
+        if (endpoint == null) {
+            return null;
+        }
+        if (endpoint.startsWith("http")) {
+            return endpoint;
+        }
+
+        // if no scheme is provided, default to https
+        logger.debug("no scheme found in endpoint [{}], defaulting to https", endpoint);
+        return "https://" + endpoint;
+    }
+
     static ProxyConfiguration buildProxyConfiguration(Logger logger, Ec2ClientSettings clientSettings) {
         if (Strings.hasText(clientSettings.proxyHost)) {
             try {

diff --git a/plugins/discovery-ec2/src/test/java/org/opensearch/discovery/ec2/AwsEc2ServiceImplTests.java b/plugins/discovery-ec2/src/test/java/org/opensearch/discovery/ec2/AwsEc2ServiceImplTests.java
@@ -202,4 +202,24 @@ public void testAWSConfigurationWithAwsSettings() {
         assertTrue(clientOverrideConfiguration.retryPolicy().isPresent());
         assertThat(clientOverrideConfiguration.retryPolicy().get().numRetries(), is(10));
     }
+
+    public void testGetFullEndpointWithScheme() {
+        final Settings settings = Settings.builder().put("discovery.ec2.endpoint", "http://ec2.us-west-2.amazonaws.com").build();
+        Ec2ClientSettings clientSettings = Ec2ClientSettings.getClientSettings(settings);
+
+        AwsEc2ServiceImpl awsEc2ServiceImpl = new AwsEc2ServiceImpl();
+
+        String endpoint = awsEc2ServiceImpl.getFullEndpoint(clientSettings.endpoint);
+        assertEquals("http://ec2.us-west-2.amazonaws.com", endpoint);
+    }
+
+    public void testGetFullEndpointWithoutScheme() {
+        final Settings settings = Settings.builder().put("discovery.ec2.endpoint", "ec2.us-west-2.amazonaws.com").build();
+        Ec2ClientSettings clientSettings = Ec2ClientSettings.getClientSettings(settings);
+
+        AwsEc2ServiceImpl awsEc2ServiceImpl = new AwsEc2ServiceImpl();
+
+        String endpoint = awsEc2ServiceImpl.getFullEndpoint(clientSettings.endpoint);
+        assertEquals("https://ec2.us-west-2.amazonaws.com", endpoint);
+    }
 }