Add dependent plugins check during plugin load

Signed-off-by: Junqiu Lei <junqiu@amazon.com>
opensearch-project · Jan 17, 2025 · e3c084f · e3c084f
1 parent 660f577
commit e3c084f
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 4 deletions.
diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java
@@ -14,6 +14,7 @@
 import java.util.Optional;
 import java.util.function.Supplier;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.service.ClusterService;
@@ -86,6 +87,14 @@ public class NeuralSearch extends Plugin implements ActionPlugin, SearchPlugin,
     private final ScoreCombinationFactory scoreCombinationFactory = new ScoreCombinationFactory();
     public static final String EXPLANATION_RESPONSE_KEY = "explanation_response";
 
+    @VisibleForTesting
+    public static final Map<String, String> DEPENDENT_PLUGINS = Map.of(
+        "opensearch-knn",
+        "org.opensearch.knn.plugin.KNNPlugin",
+        "opensearch-ml-commons",
+        "org.opensearch.ml.client.MachineLearningClient"
+    );
+
     @Override
     public Collection<Object> createComponents(
         final Client client,
@@ -100,6 +109,7 @@ public Collection<Object> createComponents(
         final IndexNameExpressionResolver indexNameExpressionResolver,
         final Supplier<RepositoriesService> repositoriesServiceSupplier
     ) {
+        validateDependentPlugins();
         NeuralSearchClusterUtil.instance().initialize(clusterService);
         NeuralQueryBuilder.initialize(clientAccessor);
         NeuralSparseQueryBuilder.initialize(clientAccessor);
@@ -204,4 +214,16 @@ public List<SearchPlugin.SearchExtSpec<?>> getSearchExts() {
             )
         );
     }
+
+    private void validateDependentPlugins() {
+        DEPENDENT_PLUGINS.forEach((pluginName, className) -> {
+            try {
+                Class.forName(className);
+                log.info("Plugin [{}] is installed", pluginName);
+            } catch (ClassNotFoundException e) {
+                throw new IllegalStateException("Neural Search plugin requires the " + pluginName + " plugin to be installed", e);
+            }
+        });
+        log.info("All required plugins {} are installed", DEPENDENT_PLUGINS.keySet());
+    }
 }
diff --git a/src/test/java/org/opensearch/neuralsearch/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/NeuralSearchTests.java
@@ -4,13 +4,44 @@
  */
 package org.opensearch.neuralsearch;
 
-import org.opensearch.knn.common.KNNConstants;
-import org.opensearch.knn.index.engine.KNNEngine;
 import org.opensearch.test.OpenSearchTestCase;
+import static org.opensearch.neuralsearch.plugin.NeuralSearch.DEPENDENT_PLUGINS;
 
 public class NeuralSearchTests extends OpenSearchTestCase {
 
-    public void testValidateKNNDependency() {
-        assertEquals(KNNConstants.LUCENE_NAME, KNNEngine.LUCENE.getName());
+    public void testDependentPluginsMapContainsRequiredPlugins() {
+        // Verify the map contains both required plugins
+        assertTrue("Map should contain k-NN plugin", DEPENDENT_PLUGINS.containsKey("opensearch-knn"));
+        assertTrue("Map should contain ML Commons plugin", DEPENDENT_PLUGINS.containsKey("opensearch-ml-commons"));
+
+        // Verify the class names are correct
+        assertEquals(
+            "k-NN plugin should map to KNNPlugin class",
+            "org.opensearch.knn.plugin.KNNPlugin",
+            DEPENDENT_PLUGINS.get("opensearch-knn")
+        );
+        assertEquals(
+            "ML Commons plugin should map to MachineLearningClient class",
+            "org.opensearch.ml.client.MachineLearningClient",
+            DEPENDENT_PLUGINS.get("opensearch-ml-commons")
+        );
+    }
+
+    public void testDependentPluginsMapIsImmutable() {
+        // Verify the map cannot be modified
+        expectThrows(UnsupportedOperationException.class, () -> { DEPENDENT_PLUGINS.put("test-plugin", "test.class.Name"); });
+    }
+
+    public void testValidatePluginsWithMissingDependencies() {
+        // Test with non-existent class
+        IllegalStateException exception = expectThrows(IllegalStateException.class, () -> {
+            try {
+                Class.forName("non.existent.Class");
+            } catch (ClassNotFoundException e) {
+                throw new IllegalStateException("Neural Search plugin requires the test plugin to be installed", e);
+            }
+        });
+        assertTrue(exception.getMessage().contains("Neural Search plugin requires"));
+        assertTrue(exception.getCause() instanceof ClassNotFoundException);
     }
 }