Java Code Examples for org.apache.nifi.util.MockFlowFile#assertAttributeEquals()

@Test
public void testMimeTypeIsOctetStreamIfConflictingWithBinaryConcat() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_BIN_AGE, "1 sec");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_CONCAT);

    createFlowFiles(runner);

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.MIME_TYPE.key(), "application/zip");
    runner.enqueue(new byte[0], attributes);
    runner.run();

    runner.assertQueueEmpty();
    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 4);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0);
    bundle.assertContentEquals("Hello, World!".getBytes("UTF-8"));
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/octet-stream");
}
 

private void testEncryptionServiceWithClientSideKMSEncryptionStrategy(byte[] data) throws InitializationException, IOException {
    TestRunner runner = createPutEncryptionTestRunner(AmazonS3EncryptionService.STRATEGY_NAME_CSE_KMS, kmsKeyId);

    final Map<String, String> attrs = new HashMap<>();
    attrs.put("filename", "test.txt");
    runner.enqueue(data, attrs);
    runner.assertValid();
    runner.run();
    runner.assertAllFlowFilesTransferred(PutS3Object.REL_SUCCESS);

    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(PutS3Object.REL_SUCCESS);
    Assert.assertEquals(1, flowFiles.size());
    Assert.assertEquals(0, runner.getFlowFilesForRelationship(PutS3Object.REL_FAILURE).size());
    MockFlowFile putSuccess = flowFiles.get(0);
    Assert.assertEquals(putSuccess.getAttribute(PutS3Object.S3_ENCRYPTION_STRATEGY), AmazonS3EncryptionService.STRATEGY_NAME_CSE_KMS);

    MockFlowFile flowFile = fetchEncryptedFlowFile(attrs, AmazonS3EncryptionService.STRATEGY_NAME_CSE_KMS, kmsKeyId);
    flowFile.assertContentEquals(data);
    flowFile.assertAttributeEquals("x-amz-wrap-alg", "kms");
    flowFile.assertAttributeEquals(PutS3Object.S3_ENCRYPTION_STRATEGY, AmazonS3EncryptionService.STRATEGY_NAME_CSE_KMS);
}
 

@Test
public void testDynamicPropertiesToAttributes() throws IOException {
    TestRunner runner = TestRunners.newTestRunner(new GenerateFlowFile());
    runner.setProperty(GenerateFlowFile.FILE_SIZE, "1B");
    runner.setProperty(GenerateFlowFile.DATA_FORMAT, GenerateFlowFile.DATA_FORMAT_TEXT);
    runner.setProperty("plain.dynamic.property", "Plain Value");
    runner.setProperty("expression.dynamic.property", "${literal('Expression Value')}");
    runner.assertValid();

    runner.run();

    runner.assertTransferCount(GenerateFlowFile.SUCCESS, 1);
    MockFlowFile generatedFlowFile = runner.getFlowFilesForRelationship(GenerateFlowFile.SUCCESS).get(0);
    generatedFlowFile.assertAttributeEquals("plain.dynamic.property", "Plain Value");
    generatedFlowFile.assertAttributeEquals("expression.dynamic.property", "Expression Value");
}
 

@Test
public void testShouldAllowNoCaptureGroups() throws Exception {
    // Arrange
    final TestRunner testRunner = TestRunners.newTestRunner(new ExtractText());
    final String attributeKey = "regex.result";
    testRunner.setProperty(attributeKey, "(?s).*");

    // Act
    testRunner.enqueue(SAMPLE_STRING.getBytes("UTF-8"));
    testRunner.run();

    // Assert
    testRunner.assertAllFlowFilesTransferred(ExtractText.REL_MATCH, 1);
    final MockFlowFile out = testRunner.getFlowFilesForRelationship(ExtractText.REL_MATCH).get(0);

    // There is no global capture group, so only "key.0" exists
    out.assertAttributeNotExists(attributeKey);
    out.assertAttributeEquals(attributeKey + ".0", SAMPLE_STRING);
}
 

@Test
public void testTransformInputWithSortrPopulatedSpec() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new JoltTransformJSON());
    runner.setValidateExpressionUsage(false);
    runner.setProperty(JoltTransformJSON.JOLT_TRANSFORM, JoltTransformJSON.SORTR);
    runner.setProperty(JoltTransformJSON.JOLT_SPEC, "abcd");
    runner.enqueue(JSON_INPUT);
    runner.run();
    runner.assertAllFlowFilesTransferred(JoltTransformJSON.REL_SUCCESS);
    final MockFlowFile transformed = runner.getFlowFilesForRelationship(JoltTransformJSON.REL_SUCCESS).get(0);
    transformed.assertAttributeExists(CoreAttributes.MIME_TYPE.key());
    transformed.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(),"application/json");
    Object transformedJson = JsonUtils.jsonToObject(new ByteArrayInputStream(transformed.toByteArray()));
    Object compareJson = JsonUtils.jsonToObject(Files.newInputStream(Paths.get("src/test/resources/TestJoltTransformJson/sortrOutput.json")));
    String transformedJsonString = JsonUtils.toJsonString(transformedJson);
    String compareJsonString = JsonUtils.toJsonString(compareJson);
    assertTrue(compareJsonString.equals(transformedJsonString));
}
 

@Test
public void testWav() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractMediaMetadata());
    runner.setProperty(ExtractMediaMetadata.METADATA_KEY_FILTER, "");
    runner.setProperty(ExtractMediaMetadata.METADATA_KEY_PREFIX, "wav.");
    runner.assertValid();

    runner.enqueue(new File("target/test-classes/testWAV.wav").toPath());
    runner.run();

    runner.assertAllFlowFilesTransferred(ExtractMediaMetadata.SUCCESS, 1);
    runner.assertTransferCount(ExtractMediaMetadata.FAILURE, 0);

    final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ExtractMediaMetadata.SUCCESS);
    MockFlowFile flowFile0 = successFiles.get(0);
    flowFile0.assertAttributeExists("filename");
    flowFile0.assertAttributeEquals("filename", "testWAV.wav");
    flowFile0.assertAttributeExists("wav.Content-Type");
    assertTrue(flowFile0.getAttribute("wav.Content-Type").startsWith("audio/vnd.wave"));
    flowFile0.assertAttributeExists("wav.X-Parsed-By");
    assertTrue(flowFile0.getAttribute("wav.X-Parsed-By").contains("org.apache.tika.parser.DefaultParser"));
    assertTrue(flowFile0.getAttribute("wav.X-Parsed-By").contains("org.apache.tika.parser.audio.AudioParser"));
    flowFile0.assertAttributeExists("wav.encoding");
    flowFile0.assertAttributeEquals("wav.encoding", "PCM_SIGNED");
}
 

@Test
public void testIgnoreZeroCaptureGroupProperty() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new ExtractText());

    testRunner.setProperty(ExtractText.INCLUDE_CAPTURE_GROUP_ZERO, "false");

    final String attributeKey = "regex.result";

    testRunner.setProperty(attributeKey, "(?s)(.*)");

    testRunner.enqueue(SAMPLE_STRING.getBytes("UTF-8"));
    testRunner.run();

    testRunner.assertAllFlowFilesTransferred(ExtractText.REL_MATCH, 1);
    final MockFlowFile out = testRunner.getFlowFilesForRelationship(ExtractText.REL_MATCH).get(0);

    // Ensure the zero capture group is not in the resultant attributes
    out.assertAttributeNotExists(attributeKey + ".0");
    out.assertAttributeEquals(attributeKey, SAMPLE_STRING);
}
 

@Test
public void testFetchElasticsearchOnTriggerWithDocNotFound() throws IOException {
    runner = TestRunners.newTestRunner(new FetchElasticsearchHttpTestProcessor(false)); // simulate doc not found
    runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200");
    runner.setProperty(FetchElasticsearchHttp.INDEX, "doc");
    runner.setValidateExpressionUsage(true);
    runner.setProperty(FetchElasticsearchHttp.DOC_ID, "${doc_id}");

    runner.setIncomingConnection(true);
    runner.enqueue(docExample, new HashMap<String, String>() {{
        put("doc_id", "28039652140");
    }});
    runner.run(1, true, true);

    // This test generates a "document not found"
    runner.assertAllFlowFilesTransferred(FetchElasticsearchHttp.REL_NOT_FOUND, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(FetchElasticsearchHttp.REL_NOT_FOUND).get(0);
    assertNotNull(out);
    out.assertAttributeEquals("doc_id", "28039652140");
}
 

@Test
public void testSimpleBinaryConcatWithTextDelimitersHeaderOnly() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_BIN_AGE, "1 sec");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_CONCAT);
    runner.setProperty(MergeContent.DELIMITER_STRATEGY, MergeContent.DELIMITER_STRATEGY_TEXT);
    runner.setProperty(MergeContent.HEADER, "@");

    createFlowFiles(runner);
    runner.run();

    runner.assertQueueEmpty();
    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 3);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0);
    bundle.assertContentEquals("@Hello, World!".getBytes("UTF-8"));
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/plain-text");
}
 

@Test
public void testNotRecursive() throws InterruptedException {
    runner.setProperty(ListHDFS.RECURSE_SUBDIRS, "false");
    proc.fileSystem.addFileStatus(new Path("/test"), new FileStatus(1L, false, 1, 1L, 0L, 0L, create777(), "owner", "group", new Path("/test/testFile.txt")));

    proc.fileSystem.addFileStatus(new Path("/test"), new FileStatus(1L, true, 1, 1L, 0L, 0L, create777(), "owner", "group", new Path("/test/testDir")));
    proc.fileSystem.addFileStatus(new Path("/test/testDir"), new FileStatus(1L, false, 1, 1L, 0L, 0L, create777(), "owner", "group", new Path("/test/testDir/1.txt")));

    // first iteration will not pick up files because it has to instead check timestamps.
    // We must then wait long enough to ensure that the listing can be performed safely and
    // run the Processor again.
    runner.run();
    Thread.sleep(TimeUnit.NANOSECONDS.toMillis(2 * ListHDFS.LISTING_LAG_NANOS));
    runner.run();

    runner.assertAllFlowFilesTransferred(ListHDFS.REL_SUCCESS, 1);

    final MockFlowFile mff1 = runner.getFlowFilesForRelationship(ListHDFS.REL_SUCCESS).get(0);
    mff1.assertAttributeEquals("path", "/test");
    mff1.assertAttributeEquals("filename", "testFile.txt");
}
 

@Test
public void testCouchbaseFatalError() throws Exception {
    String docIdExp = "doc-c";

    Bucket bucket = mock(Bucket.class);
    CouchbaseException exception = new NotConnectedException();
    when(bucket.get(docIdExp, RawJsonDocument.class))
        .thenThrow(exception);
    setupMockBucket(bucket);

    testRunner.setProperty(DOC_ID, docIdExp);

    String inputFileDataStr = "input FlowFile data";
    byte[] inFileData = inputFileDataStr.getBytes(StandardCharsets.UTF_8);
    testRunner.enqueue(inFileData);
    testRunner.run();

    testRunner.assertTransferCount(REL_SUCCESS, 0);
    testRunner.assertTransferCount(REL_ORIGINAL, 0);
    testRunner.assertTransferCount(REL_RETRY, 1);
    testRunner.assertTransferCount(REL_FAILURE, 0);
    MockFlowFile orgFile = testRunner.getFlowFilesForRelationship(REL_RETRY).get(0);
    orgFile.assertContentEquals(inputFileDataStr);
    orgFile.assertAttributeEquals(Exception.key(), exception.getClass().getName());
}
 

@Test
public void testMatchAll() {
    final TestRunner runner = TestRunners.newTestRunner(new RouteOnAttribute());
    runner.setProperty(RouteOnAttribute.ROUTE_STRATEGY, RouteOnAttribute.ROUTE_ALL_MATCH.getValue());
    runner.setProperty("RouteA", "${a:equals('b')}");
    runner.setProperty("RouteB", "${b:equals('a')}");

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("a", "b");
    attributes.put("b", "a");
    runner.enqueue(new byte[0], attributes);

    attributes.put("b", "b");
    runner.enqueue(new byte[0], attributes);

    attributes.put("a", "a");
    attributes.put("b", "b");
    runner.enqueue(new byte[0], attributes);

    runner.enqueue(new byte[0]);

    runner.run(4);

    final List<MockFlowFile> match = runner.getFlowFilesForRelationship(RouteOnAttribute.REL_MATCH);
    final List<MockFlowFile> noMatch = runner.getFlowFilesForRelationship(RouteOnAttribute.REL_NO_MATCH);

    assertEquals(1, match.size());
    assertEquals(3, noMatch.size());

    for (final MockFlowFile ff : noMatch) {
        ff.assertAttributeEquals(RouteOnAttribute.ROUTE_ATTRIBUTE_KEY, "unmatched");
    }

    final Map<String, String> matchedAttrs = match.iterator().next().getAttributes();
    assertEquals("b", matchedAttrs.get("a"));
    assertEquals("a", matchedAttrs.get("b"));
    assertEquals("matched", matchedAttrs.get(RouteOnAttribute.ROUTE_ATTRIBUTE_KEY));
}
 

@Test
public void testInsertWithGeneratedKeys() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(PutSQL.class);
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    runner.setProperty(PutSQL.OBTAIN_GENERATED_KEYS, "true");
    runner.setProperty(PutSQL.CONNECTION_POOL, "dbcp");

    recreateTable("PERSONS_AI",createPersonsAutoId);
    runner.enqueue("INSERT INTO PERSONS_AI (NAME, CODE) VALUES ('Mark', 84)".getBytes());
    runner.run();

    runner.assertAllFlowFilesTransferred(PutSQL.REL_SUCCESS, 1);
    final MockFlowFile mff = runner.getFlowFilesForRelationship(PutSQL.REL_SUCCESS).get(0);
    mff.assertAttributeEquals("sql.generated.key", "1");

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            final ResultSet rs = stmt.executeQuery("SELECT * FROM PERSONS_AI");
            assertTrue(rs.next());
            assertEquals(1, rs.getInt(1));
            assertEquals("Mark", rs.getString(2));
            assertEquals(84, rs.getInt(3));
            assertFalse(rs.next());
        }
    }
}
 

@Test
public void testExtractionWithItemCount() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractAvroMetadata());
    runner.setProperty(ExtractAvroMetadata.COUNT_ITEMS, "true");

    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/user.avsc"));
    final ByteArrayOutputStream out = getOutputStreamWithMultipleUsers(schema, 6000); // creates 2 blocks
    runner.enqueue(out.toByteArray());
    runner.run();

    runner.assertAllFlowFilesTransferred(ExtractAvroMetadata.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(ExtractAvroMetadata.REL_SUCCESS).get(0);
    flowFile.assertAttributeEquals(ExtractAvroMetadata.ITEM_COUNT_ATTR, "6000");
}
 

@Test
public void testMultipleRuleHitsWithUseOriginalDoesntApplyDefaultsRepeatedly() throws Exception {
    final Criteria criteria = getCriteria();
    criteria.setFlowFilePolicy(FlowFilePolicy.USE_ORIGINAL);
    addRule(criteria, "rule 1", Arrays.asList(
            // conditions
            "${attribute.1:equals('value.1')}"), getMap(
                    // actions
                    "attribute.2", "value.2"));
    addRule(criteria, "rule 2", Arrays.asList(
            // conditions
            "${attribute.1:equals('value.1')}"), getMap(
                    // actions
                    "attribute.3", "value.3"));
    addRule(criteria, "rule 3", Arrays.asList(
            // conditions
            "${attribute.1:equals('value.1')}"), getMap(
                    // actions
                    "attribute.2", "value.3"));

    final TestRunner runner = TestRunners.newTestRunner(new UpdateAttribute());
    runner.setAnnotationData(serialize(criteria));
    runner.setProperty("default.attr", "${default.attr}-more-stuff");

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("attribute.1", "value.1");
    runner.enqueue(TEST_CONTENT.getBytes(StandardCharsets.UTF_8), attributes);

    runner.run();

    runner.assertAllFlowFilesTransferred(UpdateAttribute.REL_SUCCESS, 1);
    final List<MockFlowFile> result = runner.getFlowFilesForRelationship(UpdateAttribute.REL_SUCCESS);

    final MockFlowFile flowfile = result.get(0);

    // ensure the attributes are as expected
    flowfile.assertAttributeEquals("default.attr", "-more-stuff");
}
 

@Test
public void testMatchedEmptyStringAttribute() throws XPathFactoryConfigurationException, IOException {
    final TestRunner testRunner = TestRunners.newTestRunner(new EvaluateXQuery());
    testRunner.setProperty(EvaluateXQuery.DESTINATION, EvaluateXQuery.DESTINATION_ATTRIBUTE);
    testRunner.setProperty("xquery.result.exist.2", "/*:fruitbasket/*[name='none']/color/text()");

    testRunner.enqueue(XML_SNIPPET);
    testRunner.run();

    testRunner.assertAllFlowFilesTransferred(EvaluateXQuery.REL_NO_MATCH, 1);
    final MockFlowFile out = testRunner.getFlowFilesForRelationship(EvaluateXQuery.REL_NO_MATCH).get(0);

    out.assertAttributeEquals("xquery.result.exist.2", null);
    testRunner.getFlowFilesForRelationship(EvaluateXQuery.REL_NO_MATCH).get(0).assertContentEquals(XML_SNIPPET);
}
 

@Test
public void testNoRetryOnErrorAlwaysOutputResponseAndPenalize() throws Exception {
    final TestGRPCServer<DummyFlowFileService> server = new TestGRPCServer<>(DummyFlowFileService.class);

    try {
        final int port = server.start(0);

        final TestRunner runner = TestRunners.newTestRunner(InvokeGRPC.class);
        runner.setProperty(InvokeGRPC.PROP_SERVICE_HOST, TestGRPCServer.HOST);
        runner.setProperty(InvokeGRPC.PROP_SERVICE_PORT, String.valueOf(port));
        runner.setProperty(InvokeGRPC.PROP_OUTPUT_RESPONSE_REGARDLESS, "true");
        runner.setProperty(InvokeGRPC.PROP_PENALIZE_NO_RETRY, "true");

        final MockFlowFile mockFlowFile = new MockFlowFile(ERROR);
        runner.enqueue(mockFlowFile);
        runner.run();
        runner.assertTransferCount(InvokeGRPC.REL_RESPONSE, 1);
        runner.assertTransferCount(InvokeGRPC.REL_SUCCESS_REQ, 0);
        runner.assertTransferCount(InvokeGRPC.REL_RETRY, 0);
        runner.assertTransferCount(InvokeGRPC.REL_NO_RETRY, 1);
        runner.assertTransferCount(InvokeGRPC.REL_FAILURE, 0);
        runner.assertPenalizeCount(1);

        final List<MockFlowFile> noRetryFiles = runner.getFlowFilesForRelationship(InvokeGRPC.REL_NO_RETRY);
        assertThat(noRetryFiles.size(), equalTo(1));
        final MockFlowFile noRetry = noRetryFiles.get(0);
        noRetry.assertAttributeEquals(InvokeGRPC.RESPONSE_CODE, String.valueOf(FlowFileReply.ResponseCode.ERROR));
        noRetry.assertAttributeEquals(InvokeGRPC.RESPONSE_BODY, "error");
        noRetry.assertAttributeEquals(InvokeGRPC.SERVICE_HOST, TestGRPCServer.HOST);
        noRetry.assertAttributeEquals(InvokeGRPC.SERVICE_PORT, String.valueOf(port));

        final List<MockFlowFile> responseFiles = runner.getFlowFilesForRelationship(InvokeGRPC.REL_RESPONSE);
        assertThat(responseFiles.size(), equalTo(1));
        final MockFlowFile response = responseFiles.get(0);
        response.assertAttributeEquals(InvokeGRPC.RESPONSE_CODE, String.valueOf(FlowFileReply.ResponseCode.ERROR));
        response.assertAttributeEquals(InvokeGRPC.RESPONSE_BODY, "error");
        response.assertAttributeEquals(InvokeGRPC.SERVICE_HOST, TestGRPCServer.HOST);
        response.assertAttributeEquals(InvokeGRPC.SERVICE_PORT, String.valueOf(port));
    } finally {
        server.stop();
    }
}
 

@Test
public void testAvroConcatWithDifferentMetadataIgnore() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_ENTRIES, "3");
    runner.setProperty(MergeContent.MIN_ENTRIES, "3");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_AVRO);
    runner.setProperty(MergeContent.METADATA_STRATEGY, MergeContent.METADATA_STRATEGY_IGNORE);

    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/TestMergeContent/user.avsc"));

    final GenericRecord user1 = new GenericData.Record(schema);
    user1.put("name", "Alyssa");
    user1.put("favorite_number", 256);
    final Map<String, String> userMeta1 = new HashMap<String, String>() {{
        put("test_metadata1", "Test 1");
    }};

    final GenericRecord user2 = new GenericData.Record(schema);
    user2.put("name", "Ben");
    user2.put("favorite_number", 7);
    user2.put("favorite_color", "red");
    final Map<String, String> userMeta2 = new HashMap<String, String>() {{
        put("test_metadata1", "Test 2"); // Test non-matching values
    }};

    final GenericRecord user3 = new GenericData.Record(schema);
    user3.put("name", "John");
    user3.put("favorite_number", 5);
    user3.put("favorite_color", "blue");
    final Map<String, String> userMeta3 = new HashMap<String, String>() {{
        put("test_metadata1", "Test 1");
        put("test_metadata2", "Test"); // Test unique
    }};

    final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
    final ByteArrayOutputStream out1 = serializeAvroRecord(schema, user1, datumWriter, userMeta1);
    final ByteArrayOutputStream out2 = serializeAvroRecord(schema, user2, datumWriter, userMeta2);
    final ByteArrayOutputStream out3 = serializeAvroRecord(schema, user3, datumWriter, userMeta3);

    runner.enqueue(out1.toByteArray());
    runner.enqueue(out2.toByteArray());
    runner.enqueue(out3.toByteArray());

    runner.run();
    runner.assertQueueEmpty();
    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 3);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0);
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/avro-binary");

    // create a reader for the merged content
    byte[] data = runner.getContentAsByteArray(bundle);
    final Map<String, GenericRecord> users = getGenericRecordMap(data, schema, "name");

    Assert.assertEquals(3, users.size());
    Assert.assertTrue(users.containsKey("Alyssa"));
    Assert.assertTrue(users.containsKey("Ben"));
    Assert.assertTrue(users.containsKey("John"));
}
 

@Test
public void testOutputResponseRegardlessWithOutputInAttributeLarge() throws Exception {
    addHandler(new GetLargeHandler(true));
    String statusUrl = "/status/200";

    setupEndpointAndRegion();

    runner.setProperty(InvokeAWSGatewayApi.PROP_RESOURCE_NAME, "/status/200");
    createFlowFiles(runner);
    runner.setProperty(InvokeAWSGatewayApi.PROP_OUTPUT_RESPONSE_REGARDLESS, "true");
    runner.setProperty(InvokeAWSGatewayApi.PROP_PUT_OUTPUT_IN_ATTRIBUTE, "outputBody");
    runner.setProperty(InvokeAWSGatewayApi.PROP_PUT_ATTRIBUTE_MAX_LENGTH, "11");

    createFlowFiles(runner);

    runner.run();

    runner.assertTransferCount(InvokeAWSGatewayApi.REL_SUCCESS_REQ_NAME, 0);
    runner.assertTransferCount(InvokeAWSGatewayApi.REL_RESPONSE_NAME, 1);
    runner.assertTransferCount(InvokeAWSGatewayApi.REL_RETRY_NAME, 0);
    runner.assertTransferCount(InvokeAWSGatewayApi.REL_NO_RETRY_NAME, 1);
    runner.assertTransferCount(InvokeAWSGatewayApi.REL_FAILURE_NAME, 0);
    runner.assertPenalizeCount(0);

    // expected in request status.code and status.message
    // original flow file (+attributes)
    final MockFlowFile bundle = runner
        .getFlowFilesForRelationship(InvokeAWSGatewayApi.REL_NO_RETRY_NAME).get(0);
    bundle.assertContentEquals("Hello".getBytes("UTF-8"));
    bundle.assertAttributeEquals(InvokeAWSGatewayApi.STATUS_CODE, "404");
    bundle.assertAttributeEquals(InvokeAWSGatewayApi.STATUS_MESSAGE, "Not Found");
    bundle.assertAttributeEquals("outputBody", "{\"name\":\"Lo");
    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(InvokeAWSGatewayApi.REL_RESPONSE_NAME).get(0);
    bundle1.assertContentEquals(
        "{\"name\":\"Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. "
            + "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor "
            + "in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, "
            + "sunt in culpa qui officia deserunt mollit anim id est laborum.\"}");
    bundle1.assertAttributeEquals(InvokeAWSGatewayApi.STATUS_CODE, "404");
    bundle1.assertAttributeEquals(InvokeAWSGatewayApi.STATUS_MESSAGE, "Not Found");
    bundle1.assertAttributeEquals("Foo", "Bar");
    bundle1.assertAttributeEquals("Content-Type", "application/json");
}
 

@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);
    runner.assertPenalizeCount(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");
}