Skip to content

Commit

Permalink
keep both dryRun and provision parameter
Browse files Browse the repository at this point in the history
Signed-off-by: Jackie Han <[email protected]>
  • Loading branch information
jackiehanyang committed Dec 14, 2023
1 parent ac9da88 commit 17b0739
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

/**
Expand All @@ -75,26 +75,29 @@ 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
*/
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;
Expand All @@ -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();
Expand All @@ -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
Expand Down Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down Expand Up @@ -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);
Expand All @@ -222,6 +222,7 @@ public void testMaxWorkflow() {
null,
template,
false,
false,
WORKFLOW_REQUEST_TIMEOUT.get(settings),
MAX_WORKFLOWS.get(settings)
);
Expand Down Expand Up @@ -259,6 +260,7 @@ public void testFailedToCreateNewWorkflow() {
null,
template,
false,
false,
WORKFLOW_REQUEST_TIMEOUT.get(settings),
MAX_WORKFLOWS.get(settings)
);
Expand Down Expand Up @@ -296,6 +298,7 @@ public void testCreateNewWorkflow() {
null,
template,
false,
false,
WORKFLOW_REQUEST_TIMEOUT.get(settings),
MAX_WORKFLOWS.get(settings)
);
Expand Down Expand Up @@ -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();

Expand All @@ -388,6 +391,7 @@ public void testCreateWorkflow_withProvisionParam() {
null,
validTemplate,
true,
true,
WORKFLOW_REQUEST_TIMEOUT.get(settings),
MAX_WORKFLOWS.get(settings)
);
Expand Down Expand Up @@ -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")
Expand All @@ -445,6 +449,7 @@ public void testCreateWorkflow_withFailedProvision() {
null,
validTemplate,
true,
true,
WORKFLOW_REQUEST_TIMEOUT.get(settings),
MAX_WORKFLOWS.get(settings)
);
Expand Down

0 comments on commit 17b0739

Please sign in to comment.