Fixes unit tests

Signed-off-by: Darshit Chanpura <[email protected]>

Author: DarshitChanpura

src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java

@@ -376,7 +376,7 @@ public void testCreateNewWorkflow() {
 
         // Bypass putInitialStateToWorkflowState and force on response
         doAnswer(invocation -> {
-            ActionListener<IndexResponse> responseListener = invocation.getArgument(3);
+            ActionListener<IndexResponse> responseListener = invocation.getArgument(4);
             responseListener.onResponse(new IndexResponse(new ShardId(WORKFLOW_STATE_INDEX, "", 1), "1", 1L, 1L, 1L, true));
             return null;
         }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(any(), any(), any(), any(), any());
@@ -450,7 +450,7 @@ public void testCreateWithUserAndFilterOn() {
 
         // Bypass putInitialStateToWorkflowState and force on response
         doAnswer(invocation -> {
-            ActionListener<IndexResponse> responseListener = invocation.getArgument(3);
+            ActionListener<IndexResponse> responseListener = invocation.getArgument(4);
             responseListener.onResponse(new IndexResponse(new ShardId(WORKFLOW_STATE_INDEX, "", 1), "1", 1L, 1L, 1L, true));
             return null;
         }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(any(), any(), any(), any(), any());
@@ -925,7 +925,7 @@ public void testCreateWorkflow_withValidation_withProvision_Success() throws Exc
 
         // Bypass putInitialStateToWorkflowState and force on response
         doAnswer(invocation -> {
-            ActionListener<IndexResponse> responseListener = invocation.getArgument(3);
+            ActionListener<IndexResponse> responseListener = invocation.getArgument(4);
             responseListener.onResponse(new IndexResponse(new ShardId(WORKFLOW_STATE_INDEX, "", 1), "1", 1L, 1L, 1L, true));
             return null;
         }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(any(), any(), any(), any(), any());
@@ -986,7 +986,7 @@ public void testCreateWorkflow_withValidation_withWaitForCompletion_withProvisio
 
         // Bypass putInitialStateToWorkflowState and force on response
         doAnswer(invocation -> {
-            ActionListener<IndexResponse> responseListener = invocation.getArgument(3);
+            ActionListener<IndexResponse> responseListener = invocation.getArgument(4);
             responseListener.onResponse(new IndexResponse(new ShardId(WORKFLOW_STATE_INDEX, "", 1), "1", 1L, 1L, 1L, true));
             return null;
         }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(any(), any(), any(), any(), any());
@@ -1063,7 +1063,7 @@ public void testCreateWorkflow_withValidation_withWaitForCompletionTimeSetZero_w
 
         // Bypass putInitialStateToWorkflowState and force on response
         doAnswer(invocation -> {
-            ActionListener<IndexResponse> responseListener = invocation.getArgument(3);
+            ActionListener<IndexResponse> responseListener = invocation.getArgument(4);
             responseListener.onResponse(new IndexResponse(new ShardId(WORKFLOW_STATE_INDEX, "", 1), "1", 1L, 1L, 1L, true));
             return null;
         }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(any(), any(), any(), any(), any());
@@ -1140,7 +1140,7 @@ public void testCreateWorkflow_withValidation_withProvision_FailedProvisioning()
 
         // Bypass putInitialStateToWorkflowState and force on response
         doAnswer(invocation -> {
-            ActionListener<IndexResponse> responseListener = invocation.getArgument(3);
+            ActionListener<IndexResponse> responseListener = invocation.getArgument(4);
             responseListener.onResponse(new IndexResponse(new ShardId(WORKFLOW_STATE_INDEX, "", 1), "1", 1L, 1L, 1L, true));
             return null;
         }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(any(), any(), any(), any(), any());

src/test/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportActionTests.java

@@ -170,6 +170,13 @@ public void testDeprovisionWorkflow() throws Exception {
             return null;
         }).when(flowFrameworkIndicesHandler).doesTemplateExist(anyString(), any(), any(), any());
 
+        // putInitialStateToWorkflowState must invoke its ActionListener.onResponse
+        doAnswer(inv -> {
+            ActionListener<?> l = inv.getArgument(4);
+            l.onResponse(null);
+            return null;
+        }).when(flowFrameworkIndicesHandler).putInitialStateToWorkflowState(anyString(), any(), any(), any(), any());
+
         PlainActionFuture<WorkflowData> future = PlainActionFuture.newFuture();
         future.onResponse(WorkflowData.EMPTY);
         when(this.deleteConnectorStep.execute(anyString(), any(WorkflowData.class), anyMap(), anyMap(), anyMap(), nullable(String.class)))
@@ -291,6 +298,21 @@ public void testAllowDeleteRequired() throws Exception {
             "These resources require the allow_delete parameter to deprovision: [index_name test-index].",
             exceptionCaptor.getValue().getMessage()
         );
+
+        // doesTemplateExist -> false for this test
+        doAnswer(inv -> {
+            java.util.function.Consumer<Boolean> cb = inv.getArgument(2);
+            cb.accept(Boolean.FALSE);
+            return null;
+        }).when(flowFrameworkIndicesHandler).doesTemplateExist(anyString(), any(), any(), any());
+
+        // deleteFlowFrameworkSystemIndexDoc must call onResponse
+        doAnswer(inv -> {
+            ActionListener<?> l = inv.getArgument(2);
+            l.onResponse(null);
+            return null;
+        }).when(flowFrameworkIndicesHandler).deleteFlowFrameworkSystemIndexDoc(anyString(), any(), any());
+
         verify(flowFrameworkIndicesHandler, times(0)).deleteResourceFromStateIndex(
             anyString(),
             nullable(String.class),