NIFI-1362 Set mime.type attribute on response FlowFile based on InvokeHTTP response Content-Type

Signed-off-by: Aldrin Piri <aldrin@apache.org>
This commit is contained in:
Jeremy Dyer 2016-01-14 23:16:55 -05:00 committed by Aldrin Piri
parent 92062f9beb
commit ecb81ec113
2 changed files with 47 additions and 2 deletions

View File

@ -556,7 +556,7 @@ public final class InvokeHTTP extends AbstractProcessor {
throw new IllegalStateException("Status code unknown, connection hasn't been attempted.");
}
// Create a map of the status attributes that are always written to the request and reponse FlowFiles
// Create a map of the status attributes that are always written to the request and response FlowFiles
Map<String, String> statusAttributes = new HashMap<>();
statusAttributes.put(STATUS_CODE, String.valueOf(statusCode));
statusAttributes.put(STATUS_MESSAGE, statusMessage);
@ -602,7 +602,7 @@ public final class InvokeHTTP extends AbstractProcessor {
responseFlowFile = session.create();
}
// write the status attributes
// write attributes to response flowfile
responseFlowFile = session.putAllAttributes(responseFlowFile, statusAttributes);
// write the response headers as attributes
@ -612,6 +612,10 @@ public final class InvokeHTTP extends AbstractProcessor {
// transfer the message body to the payload
// can potentially be null in edge cases
if (bodyExists) {
// write content type attribute to response flowfile if it is available
if (responseBody.contentType() != null) {
responseFlowFile = session.putAttribute(responseFlowFile, CoreAttributes.MIME_TYPE.key(), responseBody.contentType().toString());
}
if (teeInputStream != null) {
responseFlowFile = session.importFrom(teeInputStream, responseFlowFile);
} else {

View File

@ -225,6 +225,47 @@ public abstract class TestInvokeHttpCommon {
bundle1.assertAttributeEquals("Content-Type", "text/plain; charset=ISO-8859-1");
}
@Test
public void testOutputResponseSetMimeTypeToResponseContentType() throws Exception {
addHandler(new GetOrHeadHandler());
String statusUrl = "/status/200";
runner.setProperty(InvokeHTTP.PROP_URL, url + statusUrl);
runner.setProperty(InvokeHTTP.PROP_METHOD, "GET");
runner.setProperty(InvokeHTTP.PROP_OUTPUT_RESPONSE_REGARDLESS,"true");
runner.setProperty(InvokeHTTP.PROP_PUT_OUTPUT_IN_ATTRIBUTE,"outputBody");
createFlowFiles(runner);
runner.run();
runner.assertTransferCount(InvokeHTTP.REL_SUCCESS_REQ, 1);
runner.assertTransferCount(InvokeHTTP.REL_RESPONSE, 1);
runner.assertTransferCount(InvokeHTTP.REL_RETRY, 0);
runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY,0);
runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
// expected in request status.code and status.message
// original flow file (+attributes)
final MockFlowFile bundle = runner.getFlowFilesForRelationship(InvokeHTTP.REL_SUCCESS_REQ).get(0);
bundle.assertContentEquals("Hello".getBytes("UTF-8"));
bundle.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
bundle.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
bundle.assertAttributeEquals("outputBody", statusUrl);
bundle.assertAttributeEquals("Foo", "Bar");
// expected in response
// status code, status message, all headers from server response --> ff attributes
// server response message body into payload of ff
final MockFlowFile bundle1 = runner.getFlowFilesForRelationship(InvokeHTTP.REL_RESPONSE).get(0);
bundle1.assertContentEquals(statusUrl.getBytes("UTF-8"));
bundle1.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
bundle1.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
bundle1.assertAttributeEquals("Foo", "Bar");
bundle1.assertAttributeEquals("Content-Type", "text/plain; charset=ISO-8859-1");
bundle1.assertAttributeEquals("mime.type", "text/plain; charset=ISO-8859-1");
}
@Test
public void testOutputResponseRegardlessWithOutputInAttributeLarge() throws Exception {
addHandler(new GetLargeHandler());