From 17b0739964cf30a32a2499d2a91b79595274c5b7 Mon Sep 17 00:00:00 2001
From: Jackie Han <jkhanjob@gmail.com>
Date: Wed, 13 Dec 2023 16:27:13 -0800
Subject: [PATCH] keep both dryRun and provision parameter

Signed-off-by: Jackie Han <jkhanjob@gmail.com>
---
 .../rest/RestCreateWorkflowAction.java          |  4 +++-
 .../CreateWorkflowTransportAction.java          |  2 +-
 .../transport/WorkflowRequest.java              | 17 +++++++++++++++--
 .../CreateWorkflowTransportActionTests.java     | 15 ++++++++++-----
 4 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java
index 8fb354fa1..deeabdd76 100644
--- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java
+++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java
@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Locale;
 
+import static org.opensearch.flowframework.common.CommonValue.DRY_RUN;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
@@ -90,9 +91,10 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
 
             String workflowId = request.param(WORKFLOW_ID);
             Template template = Template.parse(request.content().utf8ToString());
+            boolean dryRun = request.paramAsBoolean(DRY_RUN, false);
             boolean provision = request.paramAsBoolean(PROVISION_WORKFLOW, false);
 
-            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template, provision, requestTimeout, maxWorkflows);
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template, dryRun, provision, requestTimeout, maxWorkflows);
 
             return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
                 XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
index de3c6fc13..765c9cae5 100644
--- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
+++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -95,7 +95,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
             user
         );
 
-        if (request.isProvision()) {
+        if (request.isDryRun()) {
             try {
                 validateWorkflows(templateWithUser);
             } catch (Exception e) {
diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java
index bc9fce9b6..057f13d01 100644
--- a/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java
+++ b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java
@@ -59,7 +59,7 @@ public class WorkflowRequest extends ActionRequest {
      * @param template the use case template which describes the workflow
      */
     public WorkflowRequest(@Nullable String workflowId, @Nullable Template template) {
-        this(workflowId, template, false, null, null);
+        this(workflowId, template, false, false, null, null);
     }
 
     /**
@@ -75,13 +75,14 @@ public WorkflowRequest(
         @Nullable TimeValue requestTimeout,
         @Nullable Integer maxWorkflows
     ) {
-        this(workflowId, template, false, requestTimeout, maxWorkflows);
+        this(workflowId, template, false, false, requestTimeout, maxWorkflows);
     }
 
     /**
      * Instantiates a new WorkflowRequest
      * @param workflowId the documentId of the workflow
      * @param template the use case template which describes the workflow
+     * @param dryRun flag to indicate if validation is necessary
      * @param provision flag to indicate if provision is necessary
      * @param requestTimeout timeout of the request
      * @param maxWorkflows max number of workflows
@@ -89,12 +90,14 @@ public WorkflowRequest(
     public WorkflowRequest(
         @Nullable String workflowId,
         @Nullable Template template,
+        boolean dryRun,
         boolean provision,
         @Nullable TimeValue requestTimeout,
         @Nullable Integer maxWorkflows
     ) {
         this.workflowId = workflowId;
         this.template = template;
+        this.dryRun = dryRun;
         this.provision = provision;
         this.requestTimeout = requestTimeout;
         this.maxWorkflows = maxWorkflows;
@@ -110,6 +113,7 @@ public WorkflowRequest(StreamInput in) throws IOException {
         this.workflowId = in.readOptionalString();
         String templateJson = in.readOptionalString();
         this.template = templateJson == null ? null : Template.parse(templateJson);
+        this.dryRun = in.readBoolean();
         this.provision = in.readBoolean();
         this.requestTimeout = in.readOptionalTimeValue();
         this.maxWorkflows = in.readOptionalInt();
@@ -133,6 +137,14 @@ public Template getTemplate() {
         return this.template;
     }
 
+    /**
+     * Gets the dry run validation flag
+     * @return the dry run boolean
+     */
+    public boolean isDryRun() {
+        return this.dryRun;
+    }
+
     /**
      * Gets the provision flag
      * @return the provision boolean
@@ -163,6 +175,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(workflowId);
         out.writeOptionalString(template == null ? null : template.toJson());
         out.writeBoolean(dryRun);
+        out.writeBoolean(provision);
         out.writeOptionalTimeValue(requestTimeout);
         out.writeOptionalInt(maxWorkflows);
     }
diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java
index 09df66169..e03847ac6 100644
--- a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java
+++ b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java
@@ -144,12 +144,12 @@ public void setUp() throws Exception {
         );
     }
 
-    public void testDryRunValidation_Success() {
+    public void testDryRunValidation_withoutProvision_Success() {
         Template validTemplate = generateValidTemplate();
 
         @SuppressWarnings("unchecked")
         ActionListener<WorkflowResponse> listener = mock(ActionListener.class);
-        WorkflowRequest createNewWorkflow = new WorkflowRequest(null, validTemplate, true, null, null);
+        WorkflowRequest createNewWorkflow = new WorkflowRequest(null, validTemplate, true, false, null, null);
         createWorkflowTransportAction.doExecute(mock(Task.class), createNewWorkflow, listener);
     }
 
@@ -207,7 +207,7 @@ public void testDryRunValidation_Failed() {
 
         @SuppressWarnings("unchecked")
         ActionListener<WorkflowResponse> listener = mock(ActionListener.class);
-        WorkflowRequest createNewWorkflow = new WorkflowRequest(null, cyclicalTemplate, true, null, null);
+        WorkflowRequest createNewWorkflow = new WorkflowRequest(null, cyclicalTemplate, true, false, null, null);
 
         createWorkflowTransportAction.doExecute(mock(Task.class), createNewWorkflow, listener);
         ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
@@ -222,6 +222,7 @@ public void testMaxWorkflow() {
             null,
             template,
             false,
+            false,
             WORKFLOW_REQUEST_TIMEOUT.get(settings),
             MAX_WORKFLOWS.get(settings)
         );
@@ -259,6 +260,7 @@ public void testFailedToCreateNewWorkflow() {
             null,
             template,
             false,
+            false,
             WORKFLOW_REQUEST_TIMEOUT.get(settings),
             MAX_WORKFLOWS.get(settings)
         );
@@ -296,6 +298,7 @@ public void testCreateNewWorkflow() {
             null,
             template,
             false,
+            false,
             WORKFLOW_REQUEST_TIMEOUT.get(settings),
             MAX_WORKFLOWS.get(settings)
         );
@@ -378,7 +381,7 @@ public void testUpdateWorkflow() {
         assertEquals("1", responseCaptor.getValue().getWorkflowId());
     }
 
-    public void testCreateWorkflow_withProvisionParam() {
+    public void testCreateWorkflow_withDryRun_withProvision_Success() {
 
         Template validTemplate = generateValidTemplate();
 
@@ -388,6 +391,7 @@ public void testCreateWorkflow_withProvisionParam() {
             null,
             validTemplate,
             true,
+            true,
             WORKFLOW_REQUEST_TIMEOUT.get(settings),
             MAX_WORKFLOWS.get(settings)
         );
@@ -436,7 +440,7 @@ public void testCreateWorkflow_withProvisionParam() {
         assertEquals("1", workflowResponseCaptor.getValue().getWorkflowId());
     }
 
-    public void testCreateWorkflow_withFailedProvision() {
+    public void testCreateWorkflow_withDryRun_withProvision_FailedProvisioning() {
         Template validTemplate = generateValidTemplate();
 
         @SuppressWarnings("unchecked")
@@ -445,6 +449,7 @@ public void testCreateWorkflow_withFailedProvision() {
             null,
             validTemplate,
             true,
+            true,
             WORKFLOW_REQUEST_TIMEOUT.get(settings),
             MAX_WORKFLOWS.get(settings)
         );