diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
index b05e43ed4..9b945a05c 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
@@ -127,22 +127,7 @@ public void testDeployModel() throws ExecutionException, InterruptedException, I
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.COMPLETED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
@@ -215,22 +200,7 @@ public void testDeployModelTaskFailure() throws IOException, InterruptedExceptio
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.FAILED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                "error",
-                null,
-                false
-            );
+            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.FAILED).async(false).error("error").build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java
index 061d9f8c8..6a6809d07 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java
@@ -144,22 +144,7 @@ public void testRegisterLocalCustomModelSuccess() throws Exception {
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.COMPLETED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
@@ -239,22 +224,7 @@ public void testRegisterLocalCustomModelDeployStateUpdateFail() throws Exception
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.COMPLETED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
@@ -342,22 +312,13 @@ public void testRegisterLocalCustomModelTaskFailure() {
         // Stub get ml task for failure case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.FAILED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                testErrorMessage,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder()
+                .taskId(taskId)
+                .modelId(modelId)
+                .state(MLTaskState.FAILED)
+                .error(testErrorMessage)
+                .async(false)
+                .build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java
index afe97bacb..162b97dba 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java
@@ -137,22 +137,7 @@ public void testRegisterLocalPretrainedModelSuccess() throws Exception {
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.COMPLETED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
@@ -247,22 +232,13 @@ public void testRegisterLocalPretrainedModelTaskFailure() {
         // Stub get ml task for failure case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.FAILED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                testErrorMessage,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder()
+                .taskId(taskId)
+                .modelId(modelId)
+                .state(MLTaskState.FAILED)
+                .error(testErrorMessage)
+                .async(false)
+                .build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java
index 9c35af33c..79d7bb883 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java
@@ -140,22 +140,7 @@ public void testRegisterLocalSparseEncodingModelSuccess() throws Exception {
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.COMPLETED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());
@@ -252,22 +237,13 @@ public void testRegisterLocalSparseEncodingModelTaskFailure() {
         // Stub get ml task for failure case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = new MLTask(
-                taskId,
-                modelId,
-                null,
-                null,
-                MLTaskState.FAILED,
-                null,
-                null,
-                null,
-                null,
-                null,
-                null,
-                testErrorMessage,
-                null,
-                false
-            );
+            MLTask output = MLTask.builder()
+                .taskId(taskId)
+                .modelId(modelId)
+                .state(MLTaskState.FAILED)
+                .error(testErrorMessage)
+                .async(false)
+                .build();
             actionListener.onResponse(output);
             return null;
         }).when(machineLearningNodeClient).getTask(any(), any());