Fixing class casting exception for getTask

Signed-off-by: Joshua Palis <[email protected]>
opensearch-project · Nov 10, 2023 · 828f0a3 · 828f0a3
1 parent 7e44eb2
commit 828f0a3
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java b/client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java
@@ -208,9 +208,7 @@ public void execute(FunctionName name, Input input, ActionListener<MLExecuteTask
     public void getTask(String taskId, ActionListener<MLTask> listener) {
         MLTaskGetRequest mlTaskGetRequest = MLTaskGetRequest.builder().taskId(taskId).build();
 
-        client.execute(MLTaskGetAction.INSTANCE, mlTaskGetRequest, ActionListener.wrap(response -> {
-            listener.onResponse(MLTaskGetResponse.fromActionResponse(response).getMlTask());
-        }, listener::onFailure));
+        client.execute(MLTaskGetAction.INSTANCE, mlTaskGetRequest, getMLTaskResponseActionListener(listener));
     }
 
     @Override
@@ -250,6 +248,16 @@ public void createConnector(MLCreateConnectorInput mlCreateConnectorInput, Actio
         client.execute(MLCreateConnectorAction.INSTANCE, createConnectorRequest, getMlCreateConnectorResponseActionListener(listener));
     }
 
+    private ActionListener<MLTaskGetResponse> getMLTaskResponseActionListener(ActionListener<MLTask> listener) {
+        ActionListener<MLTaskGetResponse> internalListener = ActionListener
+            .wrap(getResponse -> { listener.onResponse(getResponse.getMlTask()); }, listener::onFailure);
+        ActionListener<MLTaskGetResponse> actionListener = wrapActionListener(internalListener, response -> {
+            MLTaskGetResponse getResponse = MLTaskGetResponse.fromActionResponse(response);
+            return getResponse;
+        });
+        return actionListener;
+    }
+
     private ActionListener<MLDeployModelResponse> getMlDeployModelResponseActionListener(ActionListener<MLDeployModelResponse> listener) {
         ActionListener<MLDeployModelResponse> actionListener = wrapActionListener(listener, response -> {
             MLDeployModelResponse deployModelResponse = MLDeployModelResponse.fromActionResponse(response);