From 0dc66552d14d68ad3b67583973f45f84839b3c38 Mon Sep 17 00:00:00 2001
From: Owais Kazi <owaiskazi19@gmail.com>
Date: Wed, 4 Oct 2023 15:27:40 -0700
Subject: [PATCH] Added javadocs

Signed-off-by: Owais Kazi <owaiskazi19@gmail.com>
---
 .../flowframework/workflow/DeployModelStep.java   |  7 +++++++
 .../flowframework/workflow/GetTask.java           | 15 +++++++++++----
 .../flowframework/workflow/RegisterModelStep.java |  7 +++++++
 .../workflow/WorkflowStepFactory.java             |  1 +
 .../workflow/DeployModelStepTests.java            |  3 ++-
 5 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
index 3a3918bce..43ed8cbb1 100644
--- a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
+++ b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -20,6 +20,9 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
+/**
+ * Step to deploy a model
+ */
 public class DeployModelStep implements WorkflowStep {
     private static final Logger logger = LogManager.getLogger(DeployModelStep.class);
 
@@ -27,6 +30,10 @@ public class DeployModelStep implements WorkflowStep {
     private static final String MODEL_ID = "model_id";
     static final String NAME = "deploy_model";
 
+    /**
+     * Instantiate this class
+     * @param nodeClient client to instantiate MLClient
+     */
     public DeployModelStep(NodeClient nodeClient) {
         this.nodeClient = nodeClient;
     }
diff --git a/src/main/java/org/opensearch/flowframework/workflow/GetTask.java b/src/main/java/org/opensearch/flowframework/workflow/GetTask.java
index 9394b55ec..a3d1caa4e 100644
--- a/src/main/java/org/opensearch/flowframework/workflow/GetTask.java
+++ b/src/main/java/org/opensearch/flowframework/workflow/GetTask.java
@@ -17,6 +17,9 @@
 import org.opensearch.ml.common.MLTaskState;
 import org.opensearch.ml.common.transport.task.MLTaskGetResponse;
 
+/**
+ * Step to get modelID of a registered local model
+ */
 @SuppressForbidden(reason = "This class is for the future work of registering local model")
 public class GetTask {
 
@@ -24,11 +27,19 @@ public class GetTask {
     private MachineLearningNodeClient machineLearningNodeClient;
     private String taskId;
 
+    /**
+     * Instantiate this class
+     * @param machineLearningNodeClient client to instantiate ml-commons APIs
+     * @param taskId taskID of the model
+     */
     public GetTask(MachineLearningNodeClient machineLearningNodeClient, String taskId) {
         this.machineLearningNodeClient = machineLearningNodeClient;
         this.taskId = taskId;
     }
 
+    /**
+     * Invokes get task API of ml-commons
+     */
     public void getTask() {
 
         ActionListener<MLTask> actionListener = new ActionListener<>() {
@@ -51,8 +62,4 @@ public void onFailure(Exception e) {
 
     }
 
-    /*@Override
-    public void run() {
-        getTask();
-    }*/
 }
diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java
index 8f59ce5c4..74a811764 100644
--- a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java
+++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java
@@ -28,6 +28,9 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 
+/**
+ * Step to register a remote model
+ */
 public class RegisterModelStep implements WorkflowStep {
 
     private static final Logger logger = LogManager.getLogger(RegisterModelStep.class);
@@ -45,6 +48,10 @@ public class RegisterModelStep implements WorkflowStep {
     private static final String MODEL_FORMAT = "model_format";
     private static final String MODEL_CONFIG = "model_config";
 
+    /**
+     * Instantiate this class
+     * @param nodeClient client to instantiate MLClient
+     */
     public RegisterModelStep(NodeClient nodeClient) {
         this.nodeClient = nodeClient;
     }
diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java
index 1c962491f..44b710de7 100644
--- a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java
+++ b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java
@@ -29,6 +29,7 @@ public class WorkflowStepFactory {
      * Instantiate this class.
      *
      * @param client The OpenSearch client steps can use
+     * @param nodeClient NodeClient to execute transport calls
      */
     public WorkflowStepFactory(Client client, NodeClient nodeClient) {
         populateMap(client, nodeClient);
diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
index e6c4801a1..44f51f65c 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
@@ -28,7 +28,8 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import static org.mockito.Mockito.*;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
 
 @ThreadLeakScope(ThreadLeakScope.Scope.NONE)
 public class DeployModelStepTests extends OpenSearchTestCase {