Removing TODOs for RestAction constructors, adding basic unit tests f…

…or added methods in CreateIndexStep, GlobalContextHandler Signed-off-by: Joshua Palis <[email protected]>
joshpalis · Oct 9, 2023 · 288a8ae · 288a8ae
1 parent 6c3d3db
commit 288a8ae
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 12 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java
@@ -27,7 +27,6 @@
 
 import java.io.IOException;
 import java.util.HashMap;
-import java.util.Locale;
 import java.util.Map;
 
 import static org.opensearch.core.rest.RestStatus.INTERNAL_SERVER_ERROR;
@@ -103,11 +102,9 @@ public void putTemplateToGlobalContext(Template template, ActionListener<IndexRe
      */
     public void updateTemplate(String documentId, Template template, ActionListener<IndexResponse> listener) {
         if (!createIndexStep.doesIndexExist(GLOBAL_CONTEXT_INDEX)) {
-            String exceptionMessage = String.format(
-                Locale.ROOT,
-                "Failed to update template {}, global_context index does not exist.",
-                documentId
-            );
+            String exceptionMessage = "Failed to update template for workflow_id : "
+                + documentId
+                + ", global_context index does not exist.";
             logger.error(exceptionMessage);
             listener.onFailure(new Exception(exceptionMessage));
         } else {

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java
@@ -34,9 +34,7 @@ public class RestCreateWorkflowAction extends BaseRestHandler {
     /**
      * Intantiates a new RestCreateWorkflowAction
      */
-    public RestCreateWorkflowAction() {
-        // TODO : Pass settings and cluster service to constructor and add settings update consumer for request timeout value
-    }
+    public RestCreateWorkflowAction() {}
 
     @Override
     public String getName() {

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java
@@ -33,9 +33,7 @@ public class RestProvisionWorkflowAction extends BaseRestHandler {
     /**
      * Instantiates a new RestProvisionWorkflowAction
      */
-    public RestProvisionWorkflowAction() {
-        // TODO : Pass settings and cluster service to constructor and add settings update consumer for request timeout value
-    }
+    public RestProvisionWorkflowAction() {}
 
     @Override
     public String getName() {

diff --git a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java
@@ -109,4 +109,38 @@ public void testStoreResponseToGlobalContext() {
         assertEquals(GLOBAL_CONTEXT_INDEX, requestCaptor.getValue().index());
         assertEquals(documentId, requestCaptor.getValue().id());
     }
+
+    public void testUpdateTemplate() throws IOException {
+        Template template = mock(Template.class);
+        ActionListener<IndexResponse> listener = mock(ActionListener.class);
+        when(template.toXContent(any(XContentBuilder.class), eq(ToXContent.EMPTY_PARAMS))).thenAnswer(invocation -> {
+            XContentBuilder builder = invocation.getArgument(0);
+            return builder;
+        });
+        when(createIndexStep.doesIndexExist(any())).thenReturn(true);
+
+        globalContextHandler.updateTemplate("1", template, null);
+
+        ArgumentCaptor<IndexRequest> requestCaptor = ArgumentCaptor.forClass(IndexRequest.class);
+        verify(client, times(1)).index(requestCaptor.capture(), any());
+
+        assertEquals("1", requestCaptor.getValue().id());
+    }
+
+    public void testFailedUpdateTemplate() throws IOException {
+        Template template = mock(Template.class);
+        ActionListener<IndexResponse> listener = mock(ActionListener.class);
+        when(createIndexStep.doesIndexExist(any())).thenReturn(false);
+
+        globalContextHandler.updateTemplate("1", template, listener);
+        ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
+
+        verify(listener, times(1)).onFailure(exceptionCaptor.capture());
+
+        assertEquals(
+            "Failed to update template for workflow_id : 1, global_context index does not exist.",
+            exceptionCaptor.getValue().getMessage()
+        );
+
+    }
 }
diff --git a/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java
@@ -187,4 +187,18 @@ public void testInitIndexIfAbsent_IndexExist_returnFalse() {
         createIndexStep.initIndexIfAbsent(index, listener);
         assertTrue(indexMappingUpdated.get(index.getIndexName()).get());
     }
+
+    public void testDoesIndexExist() {
+        ClusterState mockClusterState = mock(ClusterState.class);
+        Metadata mockMetaData = mock(Metadata.class);
+        when(clusterService.state()).thenReturn(mockClusterState);
+        when(mockClusterState.metadata()).thenReturn(mockMetaData);
+
+        createIndexStep.doesIndexExist(GLOBAL_CONTEXT_INDEX);
+
+        ArgumentCaptor<String> indexExistsCaptor = ArgumentCaptor.forClass(String.class);
+        verify(mockMetaData, times(1)).hasIndex(indexExistsCaptor.capture());
+
+        assertEquals(GLOBAL_CONTEXT_INDEX, indexExistsCaptor.getValue());
+    }
 }