added test and switched to if/else

Signed-off-by: Amit Galitzky <[email protected]>
amitgalitz · Nov 20, 2023 · 3d8eb8e · 3d8eb8e
1 parent f6927db
commit 3d8eb8e
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 15 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -197,23 +197,23 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
     protected void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionListener<Boolean> internalListener) {
         if (!flowFrameworkIndicesHandler.doesIndexExist(CommonValue.GLOBAL_CONTEXT_INDEX)) {
             internalListener.onResponse(true);
-            return;
-        }
-        QueryBuilder query = QueryBuilders.matchAllQuery();
-        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(query).size(0).timeout(requestTimeOut);
+        } else {
+            QueryBuilder query = QueryBuilders.matchAllQuery();
+            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(query).size(0).timeout(requestTimeOut);
 
-        SearchRequest searchRequest = new SearchRequest(CommonValue.GLOBAL_CONTEXT_INDEX).source(searchSourceBuilder);
+            SearchRequest searchRequest = new SearchRequest(CommonValue.GLOBAL_CONTEXT_INDEX).source(searchSourceBuilder);
 
-        client.search(searchRequest, ActionListener.wrap(searchResponse -> {
-            if (searchResponse.getHits().getTotalHits().value >= maxWorkflow) {
-                internalListener.onResponse(false);
-            } else {
-                internalListener.onResponse(true);
-            }
-        }, exception -> {
-            logger.error("Unable to fetch the workflows {}", exception);
-            internalListener.onFailure(new FlowFrameworkException("Unable to fetch the workflows", RestStatus.BAD_REQUEST));
-        }));
+            client.search(searchRequest, ActionListener.wrap(searchResponse -> {
+                if (searchResponse.getHits().getTotalHits().value >= maxWorkflow) {
+                    internalListener.onResponse(false);
+                } else {
+                    internalListener.onResponse(true);
+                }
+            }, exception -> {
+                logger.error("Unable to fetch the workflows {}", exception);
+                internalListener.onFailure(new FlowFrameworkException("Unable to fetch the workflows", RestStatus.BAD_REQUEST));
+            }));
+        }
     }
 
     private void validateWorkflows(Template template) throws Exception {

diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java
@@ -34,6 +34,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import org.mockito.ArgumentCaptor;
 
@@ -198,6 +199,20 @@ public void testMaxWorkflow() {
         assertEquals(("Maximum workflows limit reached 1000"), exceptionCaptor.getValue().getMessage());
     }
 
+    public void testMaxWorkflowWithNoIndex() {
+        @SuppressWarnings("unchecked")
+        ActionListener<Boolean> listener = mock(ActionListener.class);
+        doAnswer(invocation -> {
+            ActionListener<Boolean> checkMaxWorkflowListener = invocation.getArgument(2);
+            checkMaxWorkflowListener.onResponse(true);
+            return null;
+        }).when(createWorkflowTransportAction).checkMaxWorkflows(any(TimeValue.class), anyInt(), any());
+        createWorkflowTransportAction.checkMaxWorkflows(new TimeValue(10, TimeUnit.SECONDS), 10, listener);
+        ArgumentCaptor<Boolean> booleanCaptor = ArgumentCaptor.forClass(Boolean.class);
+        verify(listener, times(1)).onResponse(booleanCaptor.capture());
+        assertTrue(booleanCaptor.getValue());
+    }
+
     public void testFailedToCreateNewWorkflow() {
         @SuppressWarnings("unchecked")
         ActionListener<WorkflowResponse> listener = mock(ActionListener.class);