add more UT to client module (#203) (#205)

Signed-off-by: Yaliang Wu <[email protected]>
(cherry picked from commit a4305a5)

Co-authored-by: Yaliang Wu <[email protected]>

Author: opensearch-trigger-bot[bot]

client/build.gradle

@@ -35,11 +35,11 @@ jacocoTestCoverageVerification {
         rule {
             limit {
                 counter = 'LINE'
-                minimum = 0.3 //TODO: add more test to increase coverage to 0.9
+                minimum = 0.8
             }
             limit {
                 counter = 'BRANCH'
-                minimum = 0.3 //TODO: add more test to increase coverage to 0.9
+                minimum = 0.8
             }
         }
     }

client/src/main/java/org/opensearch/ml/client/MachineLearningClient.java

@@ -78,24 +78,6 @@ default ActionFuture<MLOutput> train(MLInput mlInput, boolean asyncTask) {
      */
     void train(MLInput mlInput, boolean asyncTask, ActionListener<MLOutput> listener);
 
-    /**
-     * Execute function and return ActionFuture.
-     * @param input input data
-     * @return ActionFuture of output
-     */
-    default ActionFuture<Output> execute(Input input) {
-        PlainActionFuture<Output> actionFuture = PlainActionFuture.newFuture();
-        execute(input, actionFuture);
-        return actionFuture;
-    }
-
-    /**
-     * Execute function and return output in listener
-     * @param input input data
-     * @param listener action listener
-     */
-    void execute(Input input, ActionListener<Output> listener);
-
     /**
      * Get MLModel and return ActionFuture.
      * @param modelId id of the model

client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java

@@ -16,9 +16,6 @@
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.ml.common.parameter.*;
 import org.opensearch.ml.common.transport.MLTaskResponse;
-import org.opensearch.ml.common.transport.execute.MLExecuteTaskAction;
-import org.opensearch.ml.common.transport.execute.MLExecuteTaskRequest;
-import org.opensearch.ml.common.transport.execute.MLExecuteTaskResponse;
 import org.opensearch.ml.common.transport.model.MLModelGetRequest;
 import org.opensearch.ml.common.transport.model.MLModelGetResponse;
 import org.opensearch.ml.common.transport.model.MLModelGetAction;
@@ -74,17 +71,6 @@ public void train(MLInput mlInput, boolean asyncTask, ActionListener<MLOutput> l
         client.execute(MLTrainingTaskAction.INSTANCE, trainingTaskRequest, getMlPredictionTaskResponseActionListener(listener));
     }
 
-    @Override
-    public void execute(Input input, ActionListener<Output> listener) {
-        MLExecuteTaskRequest executeTaskRequest = MLExecuteTaskRequest.builder()
-                .input(input)
-                .build();
-
-        client.execute(MLExecuteTaskAction.INSTANCE, executeTaskRequest, ActionListener.wrap(response -> {
-            listener.onResponse(MLExecuteTaskResponse.fromActionResponse(response).getOutput());
-        }, listener::onFailure));
-    }
-
     @Override
     public void getModel(String modelId, ActionListener<MLModel> listener) {
         MLModelGetRequest mlModelGetRequest = MLModelGetRequest.builder()

client/src/test/java/org/opensearch/ml/client/MachineLearningClientTest.java

@@ -50,9 +50,26 @@ public class MachineLearningClientTest {
     SearchResponse searchResponse;
 
     private String modekId = "test_model_id";
+    private MLModel mlModel;
+    private MLTask mlTask;
+
     @Before
-    public void setUp() throws Exception {
+    public void setUp() {
         MockitoAnnotations.openMocks(this);
+        String taskId = "taskId";
+        String modelId = "modelId";
+        mlTask = MLTask.builder()
+                .taskId(taskId)
+                .modelId(modelId)
+                .functionName(FunctionName.KMEANS)
+                .build();
+
+        String modelContent = "test content";
+        mlModel = MLModel.builder()
+                .algorithm(FunctionName.KMEANS)
+                .name("test")
+                .content(modelContent)
+                .build();
 
         machineLearningClient = new MachineLearningClient() {
             @Override
@@ -72,27 +89,9 @@ public void train(MLInput mlInput, boolean asyncTask, ActionListener<MLOutput> l
                 listener.onResponse(MLTrainingOutput.builder().modelId(modekId).build());
             }
 
-            @Override
-            public void execute(Input input, ActionListener<Output> listener) {
-                listener.onResponse(new Output() {
-                    @Override
-                    public void writeTo(StreamOutput out) {
-
-                    }
-
-                    @Override
-                    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-                        builder.startObject();
-                        builder.field("test", "test_value");
-                        builder.endObject();
-                        return builder;
-                    }
-                });
-            }
-
             @Override
             public void getModel(String modelId, ActionListener<MLModel> listener) {
-                listener.onResponse(MLModel.builder().build());
+                listener.onResponse(mlModel);
             }
 
             @Override
@@ -107,7 +106,7 @@ public void searchModel(SearchRequest searchRequest, ActionListener<SearchRespon
 
             @Override
             public void getTask(String taskId, ActionListener<MLTask> listener) {
-                listener.onResponse(MLTask.builder().build());
+                listener.onResponse(mlTask);
             }
 
             @Override
@@ -185,4 +184,44 @@ public void train() {
                 .build();
         assertEquals(modekId, ((MLTrainingOutput)machineLearningClient.train(mlInput, false).actionGet()).getModelId());
     }
-}
+
+    @Test
+    public void trainAndPredict() {
+        MLInput mlInput = MLInput.builder()
+                .algorithm(FunctionName.KMEANS)
+                .parameters(mlParameters)
+                .inputDataset(new DataFrameInputDataset(input))
+                .build();
+        assertEquals(output, machineLearningClient.trainAndPredict(mlInput).actionGet());
+    }
+
+    @Test
+    public void getModel() {
+        assertEquals(mlModel, machineLearningClient.getModel("modelId").actionGet());
+    }
+
+    @Test
+    public void deleteModel() {
+        assertEquals(deleteResponse, machineLearningClient.deleteModel("modelId").actionGet());
+    }
+
+    @Test
+    public void searchModel() {
+        assertEquals(searchResponse, machineLearningClient.searchModel(new SearchRequest()).actionGet());
+    }
+
+    @Test
+    public void getTask() {
+        assertEquals(mlTask, machineLearningClient.getTask("taskId").actionGet());
+    }
+
+    @Test
+    public void deleteTask() {
+        assertEquals(deleteResponse, machineLearningClient.deleteTask("taskId").actionGet());
+    }
+
+    @Test
+    public void searchTask() {
+        assertEquals(searchResponse, machineLearningClient.searchTask(new SearchRequest()).actionGet());
+    }
+}