Skip to content

Commit

Permalink
Fix and Refactor Test with Builder Pattern for Issue #2825 (#859)
Browse files Browse the repository at this point in the history 
* Fix: Use @builder for constructor in test

Signed-off-by: Junwei Dai <junweid@amazon.com>

* Fix: change parameter for different test case

Signed-off-by: Junwei Dai <junweid@amazon.com>

---------

Signed-off-by: Junwei Dai <junweid@amazon.com>
Co-authored-by: Junwei Dai <junweid@amazon.com>
  • Loading branch information
@junweid62
junweid62 and Junwei Dai authored Sep 6, 2024
1 parent bbd089b commit 1b055e8
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 144 deletions.
34 changes: 2 additions & 32 deletions src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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());
Expand Down
57 changes: 9 additions & 48 deletions src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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());
Expand Down
40 changes: 8 additions & 32 deletions ...est/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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());
Expand Down
40 changes: 8 additions & 32 deletions ...java/org/opensearch/flowframework/workflow/RegisterLocalSparseEncodingModelStepTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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());
Expand Down

0 comments on commit 1b055e8

Please sign in to comment.