Fix: change parameter for different test case

Signed-off-by: Junwei Dai <[email protected]>
opensearch-project · Sep 6, 2024 · 0c7c77a · 0c7c77a
1 parent 2262131
commit 0c7c77a
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
@@ -200,7 +200,7 @@ public void testDeployModelTaskFailure() throws IOException, InterruptedExceptio
         // Stub getTask for success case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
+            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
@@ -312,7 +312,13 @@ public void testRegisterLocalCustomModelTaskFailure() {
         // Stub get ml task for failure case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
+            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/...est/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java b/...est/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java
@@ -232,7 +232,13 @@ public void testRegisterLocalPretrainedModelTaskFailure() {
         // Stub get ml task for failure case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
+            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/...java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java b/...java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java
@@ -237,7 +237,13 @@ public void testRegisterLocalSparseEncodingModelTaskFailure() {
         // Stub get ml task for failure case
         doAnswer(invocation -> {
             ActionListener<MLTask> actionListener = invocation.getArgument(1);
-            MLTask output = MLTask.builder().taskId(taskId).modelId(modelId).state(MLTaskState.COMPLETED).async(false).build();
+            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());