diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index fcc7bc53e..48401a611 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -62,22 +62,28 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - - String workflowId = request.param(WORKFLOW_ID); - Template template = Template.parse(request.content().utf8ToString()); - WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template); - return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> { - XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS); - channel.sendResponse(new BytesRestResponse(RestStatus.CREATED, builder)); - }, exception -> { - try { - FlowFrameworkException ex = (FlowFrameworkException) exception; - XContentBuilder exceptionBuilder = ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS); - channel.sendResponse(new BytesRestResponse(ex.getRestStatus(), exceptionBuilder)); - } catch (IOException e) { - logger.error("Failed to send back create workflow exception", e); - } - })); + try { + String workflowId = request.param(WORKFLOW_ID); + Template template = Template.parse(request.content().utf8ToString()); + WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template); + return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> { + XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS); + channel.sendResponse(new BytesRestResponse(RestStatus.CREATED, builder)); + }, exception -> { + try { + FlowFrameworkException ex = (FlowFrameworkException) exception; + XContentBuilder exceptionBuilder = ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS); + channel.sendResponse(new BytesRestResponse(ex.getRestStatus(), exceptionBuilder)); + } catch (IOException e) { + logger.error("Failed to send back create workflow exception", e); + } + })); + } catch (Exception e) { + FlowFrameworkException ex = new FlowFrameworkException(e.getMessage(), RestStatus.BAD_REQUEST); + return channel -> channel.sendResponse( + new BytesRestResponse(ex.getRestStatus(), ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS)) + ); + } } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java index fd07a91eb..2d98db7fc 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java @@ -11,6 +11,7 @@ import org.opensearch.Version; import org.opensearch.client.node.NodeClient; import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.model.Workflow; @@ -19,9 +20,9 @@ import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest; import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.rest.FakeRestChannel; import org.opensearch.test.rest.FakeRestRequest; -import java.io.IOException; import java.util.List; import java.util.Locale; import java.util.Map; @@ -87,13 +88,15 @@ public void testRestCreateWorkflowActionRoutes() { } - public void testInvalidCreateWorkflowRequest() throws IOException { + public void testInvalidCreateWorkflowRequest() throws Exception { RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST) .withPath(this.createWorkflowPath) .withContent(new BytesArray(invalidTemplate), MediaTypeRegistry.JSON) .build(); - IOException ex = expectThrows(IOException.class, () -> { createWorkflowRestAction.prepareRequest(request, nodeClient); }); - assertEquals("Unable to parse field [invalid] in a template object.", ex.getMessage()); + FakeRestChannel channel = new FakeRestChannel(request, false, 1); + createWorkflowRestAction.handleRequest(request, channel, nodeClient); + assertEquals(RestStatus.BAD_REQUEST, channel.capturedResponse().status()); + assertTrue(channel.capturedResponse().content().utf8ToString().contains("Unable to parse field [invalid] in a template object.")); } }