From 956e8233d571dc2313c26594b1bb34cb3a70e694 Mon Sep 17 00:00:00 2001
From: Joshua Palis <jpalis@amazon.com>
Date: Tue, 10 Oct 2023 00:13:47 +0000
Subject: [PATCH] Adding intial failure test case for the
 ProvisionWorkflowTransportAction. Still need to add tests for success case

Signed-off-by: Joshua Palis <jpalis@amazon.com>
---
 ...ProvisionWorkflowTransportActionTests.java | 81 +++++++++++++++++++
 1 file changed, 81 insertions(+)
 create mode 100644 src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java

diff --git a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java
new file mode 100644
index 000000000..e2ae8208e
--- /dev/null
+++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.flowframework.transport;
+
+import org.opensearch.action.get.GetResponse;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.client.Client;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.workflow.WorkflowProcessSorter;
+import org.opensearch.tasks.Task;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+
+import org.mockito.ArgumentCaptor;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class ProvisionWorkflowTransportActionTests extends OpenSearchTestCase {
+
+    private ThreadPool threadPool;
+    private Client client;
+    private WorkflowProcessSorter workflowProcessSorter;
+    private ProvisionWorkflowTransportAction provisionWorkflowTransportAction;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+        this.threadPool = mock(ThreadPool.class);
+        this.client = mock(Client.class);
+        this.workflowProcessSorter = mock(WorkflowProcessSorter.class);
+
+        this.provisionWorkflowTransportAction = new ProvisionWorkflowTransportAction(
+            mock(TransportService.class),
+            mock(ActionFilters.class),
+            threadPool,
+            client,
+            workflowProcessSorter
+        );
+
+        ThreadPool clientThreadPool = mock(ThreadPool.class);
+        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
+
+        when(client.threadPool()).thenReturn(clientThreadPool);
+        when(clientThreadPool.getThreadContext()).thenReturn(threadContext);
+    }
+
+    public void testProvisionWorkflow() {
+        // TODO : Add tests for success case
+    }
+
+    public void testFailedToRetrieveTemplateFromGlobalContext() {
+        ActionListener<WorkflowResponse> listener = mock(ActionListener.class);
+        WorkflowRequest request = new WorkflowRequest("1", null);
+        doAnswer(invocation -> {
+            ActionListener<GetResponse> responseListener = invocation.getArgument(1);
+            responseListener.onFailure(new Exception("Failed to retrieve template from global context."));
+            return null;
+        }).when(client).get(any(), any());
+
+        provisionWorkflowTransportAction.doExecute(mock(Task.class), request, listener);
+        ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
+
+        verify(listener, times(1)).onFailure(exceptionCaptor.capture());
+        assertEquals("Failed to retrieve template from global context.", exceptionCaptor.getValue().getMessage());
+    }
+
+}