Only send one extension info when initializing (#103)

Signed-off-by: Daniel Widdis <[email protected]> Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Aug 27, 2022 · 8953042 · 8953042
1 parent 7af4f07
commit 8953042
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 37 deletions.
diff --git a/src/main/java/org/opensearch/sdk/ExtensionsRunner.java b/src/main/java/org/opensearch/sdk/ExtensionsRunner.java
@@ -22,7 +22,6 @@
 import org.opensearch.common.network.NetworkService;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.PageCacheRecycler;
-import org.opensearch.extensions.DiscoveryExtension;
 import org.opensearch.extensions.ExtensionBooleanResponse;
 import org.opensearch.discovery.InitializeExtensionsRequest;
 import org.opensearch.discovery.InitializeExtensionsResponse;
@@ -125,24 +124,11 @@ DiscoveryNode getOpensearchNode() {
      */
     InitializeExtensionsResponse handleExtensionInitRequest(InitializeExtensionsRequest extensionInitRequest) {
         logger.info("Registering Extension Request received from OpenSearch");
-        InitializeExtensionsResponse initializeExtensionsResponse = new InitializeExtensionsResponse(extensionSettings.getExtensionName());
         opensearchNode = extensionInitRequest.getSourceNode();
-        // Fetch the unique ID
-        for (DiscoveryExtension de : extensionInitRequest.getExtensions()) {
-            if (de.getName().equals(extensionSettings.getExtensionName())) {
-                setUniqueId(de.getId());
-                break;
-            }
-        }
-        // We could avoid this check if we only send one extension in the initialize request, rather than the entire list
-        if (getUniqueId() == null) {
-            throw new IllegalArgumentException(
-                "Extension " + extensionSettings.getExtensionName() + " does not match any requested extension."
-            );
-        }
+        setUniqueId(extensionInitRequest.getExtension().getId());
         // Successfully initialized. Send the response.
         try {
-            return initializeExtensionsResponse;
+            return new InitializeExtensionsResponse(extensionSettings.getExtensionName());
         } finally {
             // After sending successful response to initialization, send the REST API
             port = opensearchNode.getAddress().getPort();

diff --git a/src/test/java/org/opensearch/sdk/TestExtensionsRunner.java b/src/test/java/org/opensearch/sdk/TestExtensionsRunner.java
@@ -11,22 +11,21 @@
 
 package org.opensearch.sdk;
 
-import static java.util.Collections.emptySet;
 import static java.util.Collections.emptyMap;
+import static java.util.Collections.emptySet;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.doNothing;
-import static org.mockito.Mockito.mock;
 
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.List;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -35,18 +34,18 @@
 import org.opensearch.common.io.stream.NamedWriteableRegistryResponse;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.transport.TransportAddress;
-import org.opensearch.extensions.DiscoveryExtension;
 import org.opensearch.discovery.InitializeExtensionsRequest;
 import org.opensearch.discovery.InitializeExtensionsResponse;
-import org.opensearch.extensions.OpenSearchRequest;
+import org.opensearch.extensions.DiscoveryExtension;
 import org.opensearch.extensions.ExtensionsOrchestrator.OpenSearchRequestType;
+import org.opensearch.extensions.OpenSearchRequest;
 import org.opensearch.sdk.handlers.ClusterSettingsResponseHandler;
 import org.opensearch.sdk.handlers.ClusterStateResponseHandler;
 import org.opensearch.sdk.handlers.LocalNodeResponseHandler;
 import org.opensearch.sdk.handlers.RegisterRestActionsResponseHandler;
 import org.opensearch.test.OpenSearchTestCase;
-import org.opensearch.transport.TransportService;
 import org.opensearch.transport.Transport;
+import org.opensearch.transport.TransportService;
 
 public class TestExtensionsRunner extends OpenSearchTestCase {
 
@@ -107,25 +106,23 @@ public void testHandleExtensionInitRequest() throws UnknownHostException {
             emptySet(),
             Version.CURRENT
         );
-        List<DiscoveryExtension> extensions = List.of(
-            new DiscoveryExtension(
-                EXTENSION_NAME,
-                "opensearch-sdk-1",
-                "",
-                "",
-                "",
-                sourceNode.getAddress(),
-                new HashMap<String, String>(),
-                null,
-                null
-            )
+        DiscoveryExtension extension = new DiscoveryExtension(
+            EXTENSION_NAME,
+            "opensearch-sdk-1",
+            "",
+            "",
+            "",
+            sourceNode.getAddress(),
+            new HashMap<String, String>(),
+            null,
+            null
         );
 
         // Set mocked transport service
         extensionsRunner.setExtensionTransportService(this.transportService);
         doNothing().when(this.transportService).connectToNode(sourceNode);
 
-        InitializeExtensionsRequest extensionInitRequest = new InitializeExtensionsRequest(sourceNode, extensions);
+        InitializeExtensionsRequest extensionInitRequest = new InitializeExtensionsRequest(sourceNode, extension);
 
         InitializeExtensionsResponse response = extensionsRunner.handleExtensionInitRequest(extensionInitRequest);
         // Test if name and unique ID are set