Fixed test failure

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Oct 10, 2023 · 119f4a7 · 119f4a7
1 parent 4766d1c
commit 119f4a7
Show file tree

Hide file tree

Showing 9 changed files with 70 additions and 25 deletions.
diff --git a/gradle.properties b/gradle.properties
@@ -0,0 +1,38 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+# Modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+#
+
+
+# Enable build caching
+org.gradle.caching=true
+org.gradle.warning.mode=none
+org.gradle.parallel=true
+# Workaround for https://github.com/diffplug/spotless/issues/834
+org.gradle.jvmargs=-Xmx3g -XX:+HeapDumpOnOutOfMemoryError -Xss2m \
+    --add-exports jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.parser=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.tree=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED
+options.forkOptions.memoryMaximumSize=2g
+
+# Disable duplicate project id detection
+# See https://docs.gradle.org/current/userguide/upgrading_version_6.html#duplicate_project_names_may_cause_publication_to_fail
+systemProp.org.gradle.dependency.duplicate.project.detection=false
+
+# Enforce the build to fail on deprecated gradle api usage
+systemProp.org.gradle.warning.mode=fail
+
+# forcing to use TLS1.2 to avoid failure in vault
+# see https://github.com/hashicorp/vault/issues/8750#issuecomment-631236121
+systemProp.jdk.tls.client.protocols=TLSv1.2
+
+# jvm args for faster test execution by default
+systemProp.tests.jvm.argline=-XX:TieredStopAtLevel=1 -XX:ReservedCodeCacheSize=64m
diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java
@@ -10,7 +10,6 @@
 
 import com.google.common.collect.ImmutableList;
 import org.opensearch.client.Client;
-import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
@@ -33,8 +32,6 @@
  */
 public class FlowFrameworkPlugin extends Plugin {
 
-    private NodeClient client;
-
     /**
      * Instantiate this plugin.
      */
@@ -54,7 +51,6 @@ public Collection<Object> createComponents(
         IndexNameExpressionResolver indexNameExpressionResolver,
         Supplier<RepositoriesService> repositoriesServiceSupplier
     ) {
-        // TODO: Creating NodeClient is a temporary fix until we get the NodeClient from the provision API
         WorkflowStepFactory workflowStepFactory = new WorkflowStepFactory(clusterService, client);
         WorkflowProcessSorter workflowProcessSorter = new WorkflowProcessSorter(workflowStepFactory, threadPool);
 

diff --git a/src/main/java/org/opensearch/flowframework/client/MLClient.java b/src/main/java/org/opensearch/flowframework/client/MLClient.java
@@ -8,7 +8,7 @@
  */
 package org.opensearch.flowframework.client;
 
-import org.opensearch.client.node.NodeClient;
+import org.opensearch.client.Client;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 
 /**
@@ -22,12 +22,12 @@ private MLClient() {}
     /**
      * Creates machine learning client.
      *
-     * @param nodeClient node client of OpenSearch.
+     * @param client client of OpenSearch.
      * @return machine learning client from ml-commons.
      */
-    public static MachineLearningNodeClient createMLClient(NodeClient nodeClient) {
+    public static MachineLearningNodeClient createMLClient(Client client) {
         if (INSTANCE == null) {
-            INSTANCE = new MachineLearningNodeClient(nodeClient);
+            INSTANCE = new MachineLearningNodeClient(client);
         }
         return INSTANCE;
     }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -10,7 +10,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.opensearch.client.node.NodeClient;
+import org.opensearch.client.Client;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.flowframework.client.MLClient;
 import org.opensearch.ml.client.MachineLearningNodeClient;
@@ -26,24 +26,24 @@
 public class DeployModelStep implements WorkflowStep {
     private static final Logger logger = LogManager.getLogger(DeployModelStep.class);
 
-    private NodeClient nodeClient;
+    private Client client;
     private static final String MODEL_ID = "model_id";
     static final String NAME = "deploy_model";
 
     /**
      * Instantiate this class
-     * @param nodeClient client to instantiate MLClient
+     * @param client client to instantiate MLClient
      */
-    public DeployModelStep(NodeClient nodeClient) {
-        this.nodeClient = nodeClient;
+    public DeployModelStep(Client client) {
+        this.client = client;
     }
 
     @Override
     public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
 
         CompletableFuture<WorkflowData> deployModelFuture = new CompletableFuture<>();
 
-        MachineLearningNodeClient machineLearningNodeClient = MLClient.createMLClient(nodeClient);
+        MachineLearningNodeClient machineLearningNodeClient = MLClient.createMLClient(client);
 
         ActionListener<MLDeployModelResponse> actionListener = new ActionListener<>() {
             @Override

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java
@@ -10,7 +10,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.opensearch.client.node.NodeClient;
+import org.opensearch.client.Client;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.flowframework.client.MLClient;
 import org.opensearch.ml.client.MachineLearningNodeClient;
@@ -35,7 +35,7 @@ public class RegisterModelStep implements WorkflowStep {
 
     private static final Logger logger = LogManager.getLogger(RegisterModelStep.class);
 
-    private NodeClient nodeClient;
+    private Client client;
 
     static final String NAME = "register_model";
 
@@ -50,18 +50,18 @@ public class RegisterModelStep implements WorkflowStep {
 
     /**
      * Instantiate this class
-     * @param nodeClient client to instantiate MLClient
+     * @param client client to instantiate MLClient
      */
-    public RegisterModelStep(NodeClient nodeClient) {
-        this.nodeClient = nodeClient;
+    public RegisterModelStep(Client client) {
+        this.client = client;
     }
 
     @Override
     public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
 
         CompletableFuture<WorkflowData> registerModelFuture = new CompletableFuture<>();
 
-        MachineLearningNodeClient machineLearningNodeClient = MLClient.createMLClient(nodeClient);
+        MachineLearningNodeClient machineLearningNodeClient = MLClient.createMLClient(client);
 
         ActionListener<MLRegisterModelResponse> actionListener = new ActionListener<>() {
             @Override

diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java
@@ -9,7 +9,6 @@
 package org.opensearch.flowframework.workflow;
 
 import org.opensearch.client.Client;
-import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.service.ClusterService;
 
 import java.util.HashMap;
@@ -40,8 +39,8 @@ public WorkflowStepFactory(ClusterService clusterService, Client client) {
     private void populateMap(ClusterService clusterService, Client client) {
         stepMap.put(CreateIndexStep.NAME, new CreateIndexStep(clusterService, client));
         stepMap.put(CreateIngestPipelineStep.NAME, new CreateIngestPipelineStep(client));
-        stepMap.put(RegisterModelStep.NAME, new RegisterModelStep((NodeClient) client));
-        stepMap.put(DeployModelStep.NAME, new DeployModelStep((NodeClient) client));
+        stepMap.put(RegisterModelStep.NAME, new RegisterModelStep(client));
+        stepMap.put(DeployModelStep.NAME, new DeployModelStep(client));
 
         // TODO: These are from the demo class as placeholders, remove when demos are deleted
         stepMap.put("demo_delay_3", new DemoWorkflowStep(3000));

diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java
@@ -10,6 +10,8 @@
 
 import org.opensearch.client.AdminClient;
 import org.opensearch.client.Client;
+import org.opensearch.client.ClusterAdminClient;
+import org.opensearch.client.node.NodeClient;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.TestThreadPool;
 import org.opensearch.threadpool.ThreadPool;
@@ -23,13 +25,21 @@
 public class FlowFrameworkPluginTests extends OpenSearchTestCase {
 
     private Client client;
+    private NodeClient nodeClient;
+
+    private AdminClient adminClient;
+
+    private ClusterAdminClient clusterAdminClient;
     private ThreadPool threadPool;
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
         client = mock(Client.class);
-        when(client.admin()).thenReturn(mock(AdminClient.class));
+        adminClient = mock(AdminClient.class);
+        clusterAdminClient = mock(ClusterAdminClient.class);
+        when(client.admin()).thenReturn(adminClient);
+        when(adminClient.cluster()).thenReturn(clusterAdminClient);
         threadPool = new TestThreadPool(FlowFrameworkPluginTests.class.getName());
     }
 

diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
@@ -30,6 +30,7 @@
 
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
 
 @ThreadLeakScope(ThreadLeakScope.Scope.NONE)
 public class DeployModelStepTests extends OpenSearchTestCase {
@@ -74,7 +75,7 @@ public void testDeployModel() {
         CompletableFuture<WorkflowData> future = deployModel.execute(List.of(inputData));
 
         // TODO: Find a way to verify the below
-        // verify(machineLearningNodeClient).deploy(eq(MLRegisterModelInput.class), actionListenerCaptor.capture());
+        //verify(machineLearningNodeClient).deploy(eq("modelId"), actionListenerCaptor.capture());
 
         assertTrue(future.isCompletedExceptionally());
 

diff --git a/src/test/java/org/opensearch/flowframework/workflow/WorkflowProcessSorterTests.java b/src/test/java/org/opensearch/flowframework/workflow/WorkflowProcessSorterTests.java
@@ -60,6 +60,7 @@ public static void setup() {
         AdminClient adminClient = mock(AdminClient.class);
         ClusterService clusterService = mock(ClusterService.class);
         Client client = mock(Client.class);
+
         when(client.admin()).thenReturn(adminClient);
 
         testThreadPool = new TestThreadPool(WorkflowProcessSorterTests.class.getName());