From ad2a9897c30b885506fe0f6fa0cdf165fc65be3e Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Sun, 7 Jan 2024 15:06:35 -0800
Subject: [PATCH] Improve test coverage

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../CreateWorkflowTransportAction.java        |  2 +-
 .../CreateWorkflowTransportActionTests.java   | 22 +++++++++++++++----
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
index fa164c198..2fd9bd042 100644
--- 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
index 9aac97b14..ce08cdb8a 100644
--- 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,6 +226,8 @@ public void testMaxWorkflow() {
     }
 
     public void testMaxWorkflowWithNoIndex() {
+        when(flowFrameworkIndicesHandler.doesIndexExist(anyString())).thenReturn(false);
+
         ActionListener<Boolean> listener = new ActionListener<Boolean>() {
             @Override
             public void onResponse(Boolean booleanResponse) {
@@ -223,7 +235,9 @@ public void onResponse(Boolean 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);
     }