Fixed test

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Feb 17, 2024 · 666d901 · 666d901
1 parent 472bf22
commit 666d901
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 7 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -87,11 +87,7 @@ public void onResponse(MLDeployModelResponse mlDeployModelResponse) {
                         String resourceName = getResourceByWorkflowStep(getName());
                         String id = getResourceId(mlTask);
                         deployModelFuture.onResponse(
-                            new WorkflowData(
-                                Map.ofEntries(Map.entry(resourceName, id), Map.entry(MODEL_ID, mlTask.getModelId())),
-                                currentNodeInputs.getWorkflowId(),
-                                currentNodeId
-                            )
+                            new WorkflowData(Map.ofEntries(Map.entry(resourceName, id)), currentNodeInputs.getWorkflowId(), currentNodeId)
                         );
                     }, e -> { deployModelFuture.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e))); })
                 );

diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeployModelStepTests.java
@@ -45,7 +45,6 @@
 import static org.opensearch.action.DocWriteResponse.Result.UPDATED;
 import static org.opensearch.flowframework.common.CommonValue.FLOW_FRAMEWORK_THREAD_POOL_PREFIX;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW_THREAD_POOL;
-import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_STATE_INDEX;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_THREAD_POOL;
 import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
@@ -171,7 +170,6 @@ public void testDeployModel() throws ExecutionException, InterruptedException, I
         verify(machineLearningNodeClient, times(1)).getTask(any(), any());
 
         assertEquals(modelId, future.get().getContent().get(MODEL_ID));
-        assertEquals(status, future.get().getContent().get(REGISTER_MODEL_STATUS));
     }
 
     public void testDeployModelFailure() {