Improve test coverage

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Jan 8, 2024 · ad2a989 · ad2a989
1 parent afaa06d
commit ad2a989
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -257,7 +257,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
      *  @param maxWorkflow max workflows
      *  @param internalListener listener for search request
      */
-    protected void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionListener<Boolean> internalListener) {
+    void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionListener<Boolean> internalListener) {
         if (!flowFrameworkIndicesHandler.doesIndexExist(CommonValue.GLOBAL_CONTEXT_INDEX)) {
             internalListener.onResponse(true);
         } else {

diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java
@@ -8,8 +8,11 @@
  */
 package org.opensearch.flowframework.transport;
 
+import org.apache.lucene.search.TotalHits;
 import org.opensearch.Version;
 import org.opensearch.action.index.IndexResponse;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.update.UpdateResponse;
 import org.opensearch.client.Client;
@@ -27,6 +30,8 @@
 import org.opensearch.flowframework.model.WorkflowNode;
 import org.opensearch.flowframework.workflow.WorkflowProcessSorter;
 import org.opensearch.plugins.PluginsService;
+import org.opensearch.search.SearchHit;
+import org.opensearch.search.SearchHits;
 import org.opensearch.tasks.Task;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -199,15 +204,20 @@ public void testValidation_Failed() throws Exception {
     }
 
     public void testMaxWorkflow() {
+        when(flowFrameworkIndicesHandler.doesIndexExist(anyString())).thenReturn(true);
+
         @SuppressWarnings("unchecked")
         ActionListener<WorkflowResponse> listener = mock(ActionListener.class);
         WorkflowRequest workflowRequest = new WorkflowRequest(null, template, new String[] { "off" }, false);
 
         doAnswer(invocation -> {
-            ActionListener<Boolean> checkMaxWorkflowListener = invocation.getArgument(2);
-            checkMaxWorkflowListener.onResponse(false);
+            ActionListener<SearchResponse> searchListener = invocation.getArgument(1);
+            SearchResponse searchResponse = mock(SearchResponse.class);
+            SearchHits searchHits = new SearchHits(new SearchHit[0], new TotalHits(3, TotalHits.Relation.EQUAL_TO), 1.0f);
+            when(searchResponse.getHits()).thenReturn(searchHits);
+            searchListener.onResponse(searchResponse);
             return null;
-        }).when(createWorkflowTransportAction).checkMaxWorkflows(any(TimeValue.class), anyInt(), any());
+        }).when(client).search(any(SearchRequest.class), any());
 
         createWorkflowTransportAction.doExecute(mock(Task.class), workflowRequest, listener);
         ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
@@ -216,14 +226,18 @@ public void testMaxWorkflow() {
     }
 
     public void testMaxWorkflowWithNoIndex() {
+        when(flowFrameworkIndicesHandler.doesIndexExist(anyString())).thenReturn(false);
+
         ActionListener<Boolean> listener = new ActionListener<Boolean>() {
             @Override
             public void onResponse(Boolean booleanResponse) {
                 assertTrue(booleanResponse);
             }
 
             @Override
-            public void onFailure(Exception e) {}
+            public void onFailure(Exception e) {
+                fail("Should call onResponse");
+            }
         };
         createWorkflowTransportAction.checkMaxWorkflows(new TimeValue(10, TimeUnit.SECONDS), 10, listener);
     }