Java Code Examples for org.apache.nifi.util.TestRunner#assertQueueEmpty()

/**
 * Test if the per-flowfile schema have priority on the processor-wide one
 * @throws IOException
 */
@Test
public void onTriggerUsePerFlowfileSchemaIfAvailable() throws IOException {
    TestRunner runner = TestRunners.newTestRunner(new ProtobufEncoder());
    runner.setProperty("protobuf.schemaPath", ProtobufEncoderTest.class.getResource("/schemas/AddressBook.desc").getPath());

    InputStream jsonFile = ProtobufEncoderTest.class.getResourceAsStream("/data/Person.json");
    HashMap<String, String> personProperties = new HashMap<>();
    personProperties.put("protobuf.schemaPath", ProtobufEncoderTest.class.getResource("/schemas/Person.desc").getPath());
    personProperties.put("protobuf.messageType", "Person");
    runner.enqueue(jsonFile, personProperties);

    runner.assertValid();
    runner.run(1);
    runner.assertQueueEmpty();

    runner.assertAllFlowFilesTransferred(ProtobufEncoder.SUCCESS);
    List<MockFlowFile> results = runner.getFlowFilesForRelationship(ProtobufEncoder.SUCCESS);
    Assert.assertEquals("The encoder should use the schema from flowfile instead of processor if given", 1, results.size());
    results.get(0).assertContentEquals(ProtobufEncoderTest.class.getResourceAsStream("/data/Person.data"));
}
 

@Test
public void testNoSplitterInString() {

    String content = "UVAT";

    final TestRunner runner = TestRunners.newTestRunner(new SplitContent());
    runner.setProperty(SplitContent.FORMAT, SplitContent.UTF8_FORMAT.getValue());
    runner.setProperty(SplitContent.BYTE_SEQUENCE, ",");
    runner.setProperty(SplitContent.KEEP_SEQUENCE, "false");
    runner.setProperty(SplitContent.BYTE_SEQUENCE_LOCATION, SplitContent.TRAILING_POSITION.getValue());

    runner.enqueue(content.getBytes());
    runner.run();

    runner.assertTransferCount(SplitContent.REL_SPLITS, 1);
    MockFlowFile splitResult = runner.getFlowFilesForRelationship(SplitContent.REL_SPLITS).get(0);
    splitResult.assertAttributeExists(FRAGMENT_ID);
    splitResult.assertAttributeExists(SEGMENT_ORIGINAL_FILENAME);
    splitResult.assertAttributeEquals(FRAGMENT_COUNT, "1");
    splitResult.assertAttributeEquals(FRAGMENT_INDEX, "1");
    runner.assertTransferCount(SplitContent.REL_ORIGINAL, 1);

    runner.assertQueueEmpty();
    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(SplitContent.REL_SPLITS);
    splits.get(0).assertContentEquals(content);
}
 

@Test
public void testBase32RoundTrip() throws IOException {
    final TestRunner testRunner = TestRunners.newTestRunner(new EncodeContent());

    testRunner.setProperty(EncodeContent.MODE, EncodeContent.ENCODE_MODE);
    testRunner.setProperty(EncodeContent.ENCODING, EncodeContent.BASE32_ENCODING);

    testRunner.enqueue(Paths.get("src/test/resources/hello.txt"));
    testRunner.clearTransferState();
    testRunner.run();

    testRunner.assertAllFlowFilesTransferred(EncodeContent.REL_SUCCESS, 1);

    MockFlowFile flowFile = testRunner.getFlowFilesForRelationship(EncodeContent.REL_SUCCESS).get(0);
    testRunner.assertQueueEmpty();

    testRunner.setProperty(EncodeContent.MODE, EncodeContent.DECODE_MODE);
    testRunner.enqueue(flowFile);
    testRunner.clearTransferState();
    testRunner.run();
    testRunner.assertAllFlowFilesTransferred(EncodeContent.REL_SUCCESS, 1);

    flowFile = testRunner.getFlowFilesForRelationship(EncodeContent.REL_SUCCESS).get(0);
    flowFile.assertContentEquals(new File("src/test/resources/hello.txt"));
}
 

@Test
public void testSimpleBinaryConcatWithTextDelimiters() 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, "@");
    runner.setProperty(MergeContent.DEMARCATOR, "#");
    runner.setProperty(MergeContent.FOOTER, "$");

    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 testDefault() throws IOException, InitializationException {
    XMLRecordSetWriter writer = new XMLRecordSetWriter();
    TestRunner runner = setup(writer);

    runner.setProperty(writer, XMLRecordSetWriter.ROOT_TAG_NAME, "root");

    runner.enableControllerService(writer);
    runner.enqueue("");
    runner.run();
    runner.assertQueueEmpty();
    runner.assertAllFlowFilesTransferred(TestXMLRecordSetWriterProcessor.SUCCESS, 1);

    String expected = "<root><array_record><array_field>1</array_field><array_field></array_field><array_field>3</array_field>" +
            "<name1>val1</name1><name2></name2></array_record>" +
            "<array_record><array_field>1</array_field><array_field></array_field><array_field>3</array_field>" +
            "<name1>val1</name1><name2></name2></array_record></root>";
    String actual = new String(runner.getContentAsByteArray(runner.getFlowFilesForRelationship(TestXMLRecordSetWriterProcessor.SUCCESS).get(0)));
    assertThat(expected, CompareMatcher.isSimilarTo(actual).ignoreWhitespace().withNodeMatcher(new DefaultNodeMatcher(ElementSelectors.byNameAndText)));
}
 

@Test
public void testInitialDateFilter() throws IOException, SolrServerException, ParseException {
    final Date dateToFilter = DATE_FORMAT.parse(DATE_STRING_LATER);
    final org.apache.nifi.processors.solr.TestGetSolr.TestableProcessor proc = new org.apache.nifi.processors.solr.TestGetSolr.TestableProcessor(solrClient);

    TestRunner runner = createDefaultTestRunner(proc);
    runner.setProperty(GetSolr.DATE_FILTER, DATE_FORMAT.format(dateToFilter));
    runner.setProperty(GetSolr.BATCH_SIZE, "1");

    SolrInputDocument doc10 = new SolrInputDocument();
    doc10.addField("id", "doc10");
    doc10.addField("created", dateToFilter);
    SolrInputDocument doc11 = new SolrInputDocument();
    doc11.addField("id", "doc11");
    doc11.addField("created", dateToFilter);

    solrClient.add(doc10);
    solrClient.add(doc11);
    solrClient.commit();

    runner.run(1,true, true);
    runner.assertQueueEmpty();
    runner.assertAllFlowFilesTransferred(GetSolr.REL_SUCCESS, 2);
    runner.assertAllFlowFilesContainAttribute(CoreAttributes.MIME_TYPE.key());
}
 

@Test
public void testRoundTrip() throws IOException {
    final TestRunner testRunner = TestRunners.newTestRunner(new Base64EncodeContent());

    testRunner.setProperty(Base64EncodeContent.MODE, Base64EncodeContent.ENCODE_MODE);

    testRunner.enqueue(Paths.get("src/test/resources/hello.txt"));
    testRunner.clearTransferState();
    testRunner.run();

    testRunner.assertAllFlowFilesTransferred(Base64EncodeContent.REL_SUCCESS, 1);

    MockFlowFile flowFile = testRunner.getFlowFilesForRelationship(Base64EncodeContent.REL_SUCCESS).get(0);
    testRunner.assertQueueEmpty();

    testRunner.setProperty(Base64EncodeContent.MODE, Base64EncodeContent.DECODE_MODE);
    testRunner.enqueue(flowFile);
    testRunner.clearTransferState();
    testRunner.run();
    testRunner.assertAllFlowFilesTransferred(Base64EncodeContent.REL_SUCCESS, 1);

    flowFile = testRunner.getFlowFilesForRelationship(Base64EncodeContent.REL_SUCCESS).get(0);
    flowFile.assertContentEquals(new File("src/test/resources/hello.txt"));
}
 

@Test
public void testWithSingleByteSplit() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitContent());
    runner.setProperty(SplitContent.KEEP_SEQUENCE, "false");
    runner.setProperty(SplitContent.BYTE_SEQUENCE.getName(), "FF");

    runner.enqueue(new byte[]{1, 2, 3, 4, 5, (byte) 0xFF, 5, 4, 3, 2, 1});
    runner.run();

    runner.assertTransferCount(SplitContent.REL_ORIGINAL, 1);
    runner.getFlowFilesForRelationship(SplitContent.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT, "2");
    runner.assertTransferCount(SplitContent.REL_SPLITS, 2);

    runner.assertQueueEmpty();
    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(SplitContent.REL_SPLITS);
    final MockFlowFile split1 = splits.get(0);
    final MockFlowFile split2 = splits.get(1);

    split1.assertContentEquals(new byte[]{1, 2, 3, 4, 5});
    split2.assertContentEquals(new byte[]{5, 4, 3, 2, 1});
}
 

@Test
public void testArrayWrapping() throws IOException, InitializationException {
    XMLRecordSetWriter writer = new XMLRecordSetWriter();
    TestRunner runner = setup(writer);

    runner.setProperty(writer, XMLRecordSetWriter.ROOT_TAG_NAME, "root");
    runner.setProperty(writer, XMLRecordSetWriter.RECORD_TAG_NAME, "record");

    runner.setProperty(writer, XMLRecordSetWriter.ARRAY_WRAPPING, XMLRecordSetWriter.USE_PROPERTY_AS_WRAPPER);
    runner.setProperty(writer, XMLRecordSetWriter.ARRAY_TAG_NAME, "wrap");

    runner.enableControllerService(writer);
    runner.enqueue("");
    runner.run();
    runner.assertQueueEmpty();
    runner.assertAllFlowFilesTransferred(TestXMLRecordSetWriterProcessor.SUCCESS, 1);

    String expected = "<root><record><wrap><array_field>1</array_field><array_field></array_field><array_field>3</array_field></wrap>" +
            "<name1>val1</name1><name2></name2></record>" +
            "<record><wrap><array_field>1</array_field><array_field></array_field><array_field>3</array_field></wrap>" +
            "<name1>val1</name1><name2></name2></record></root>";
    String actual = new String(runner.getContentAsByteArray(runner.getFlowFilesForRelationship(TestXMLRecordSetWriterProcessor.SUCCESS).get(0)));
    assertThat(expected, CompareMatcher.isSimilarTo(actual).ignoreWhitespace().withNodeMatcher(new DefaultNodeMatcher(ElementSelectors.byNameAndText)));
}
 

@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 testSimpleBinaryConcatWithTextDelimiters() 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, "@");
    runner.setProperty(MergeContent.DEMARCATOR, "#");
    runner.setProperty(MergeContent.FOOTER, "$");

    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");
}
 

/**
 * NIFI-4326 adds a new feature to disable strict address parsing for
 * mailbox list header fields. This is a test case that asserts that
 * lax address parsing passes (when set to "strict=false") for malformed
 * addresses.
 */
@Test
public void testNonStrictParsingPassesForInvalidAddresses() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailHeaders());
    runner.setProperty(ExtractEmailHeaders.STRICT_PARSING, "false");

    MimeMessage simpleEmailMimeMessage = attachmentGenerator.SimpleEmailMimeMessage();

    simpleEmailMimeMessage.setHeader("From", "<bad_email>");
    simpleEmailMimeMessage.setHeader("To", "<>, Joe, \"\" <>");

    ByteArrayOutputStream messageBytes = new ByteArrayOutputStream();
    try {
        simpleEmailMimeMessage.writeTo(messageBytes);
    } catch (IOException | MessagingException e) {
        e.printStackTrace();
    }

    runner.enqueue(messageBytes.toByteArray());
    runner.run();

    runner.assertTransferCount(ExtractEmailHeaders.REL_SUCCESS, 1);
    runner.assertTransferCount(ExtractEmailHeaders.REL_FAILURE, 0);


    runner.assertQueueEmpty();
    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(ExtractEmailHeaders.REL_SUCCESS);
    splits.get(0).assertAttributeEquals("email.headers.from.0", "bad_email");
    splits.get(0).assertAttributeEquals("email.headers.to.0", "");
    splits.get(0).assertAttributeEquals("email.headers.to.1", "Joe");
    splits.get(0).assertAttributeEquals("email.headers.to.2", "");
}
 

/**
 * Test decoding valid files given a .desc schema
 * @throws IOException
 */
@Test
public void onTriggerDecodeValidFiles() throws IOException {
    TestRunner runner = TestRunners.newTestRunner(new ProtobufDecoder());

    // AddressBook test
    HashMap<String, String> addressBookProperties = new HashMap<>();
    addressBookProperties.put("protobuf.schemaPath", ProtobufDecoderTest.class.getResource("/schemas/AddressBook.desc").getPath());
    addressBookProperties.put("protobuf.messageType", "AddressBook");

    // AddressBook test
    for (String filename: validTestFiles) {
        InputStream jsonFile = ProtobufDecoderTest.class.getResourceAsStream("/data/" + filename + ".data");
        addressBookProperties.put("testfile", filename);
        runner.enqueue(jsonFile, addressBookProperties);
    }

    // Ensure the configuration is valid as-is
    runner.assertValid();

    // Run the enqueued content, it also takes an int = number of contents queued
    runner.run(validTestFiles.length);
    runner.assertQueueEmpty();

    // Check if the data was processed without failure
    List<MockFlowFile> results = runner.getFlowFilesForRelationship(ProtobufDecoder.SUCCESS);
    Assert.assertEquals("All flowfiles should be returned to success", validTestFiles.length, results.size());

    // Check if the content of the flowfile is as expected
    ObjectMapper mapper = new ObjectMapper();

    for (MockFlowFile result: results) {
        JsonNode expected = mapper.readTree(this.getClass().getResourceAsStream("/data/" + result.getAttribute("testfile") + ".json"));
        JsonNode given = mapper.readTree(runner.getContentAsByteArray(result));
        Assert.assertEquals("The parsing result of " + result.getAttribute("testfile") + ".data is not as expected", expected, given);
    }
}
 

/**
 * Test case added for NIFI-4326 for a potential NPE bug
 * if the email message contains no recipient header fields, ie,
 * TO, CC, BCC.
 */
@Test
public void testValidEmailWithNoRecipients() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailHeaders());
    runner.setProperty(ExtractEmailHeaders.CAPTURED_HEADERS, "MIME-Version");

    MimeMessage simpleEmailMimeMessage = attachmentGenerator.SimpleEmailMimeMessage();

    simpleEmailMimeMessage.removeHeader("To");
    simpleEmailMimeMessage.removeHeader("Cc");
    simpleEmailMimeMessage.removeHeader("Bcc");

    ByteArrayOutputStream messageBytes = new ByteArrayOutputStream();
    try {
        simpleEmailMimeMessage.writeTo(messageBytes);
    } catch (IOException | MessagingException e) {
        e.printStackTrace();
    }

    runner.enqueue(messageBytes.toByteArray());
    runner.run();

    runner.assertTransferCount(ExtractEmailHeaders.REL_SUCCESS, 1);
    runner.assertTransferCount(ExtractEmailHeaders.REL_FAILURE, 0);

    runner.assertQueueEmpty();
    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(ExtractEmailHeaders.REL_SUCCESS);
    splits.get(0).assertAttributeEquals("email.headers.from.0", from);
    splits.get(0).assertAttributeExists("email.headers.mime-version");
    splits.get(0).assertAttributeNotExists("email.headers.to");
    splits.get(0).assertAttributeNotExists("email.headers.cc");
    splits.get(0).assertAttributeNotExists("email.headers.bcc");
}
 

@Test
public void testSimpleBinaryConcatWithFileDelimiters() 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_FILENAME);
    runner.setProperty(MergeContent.HEADER, "${header}");
    runner.setProperty(MergeContent.DEMARCATOR, "${demarcator}");
    runner.setProperty(MergeContent.FOOTER, "${footer}");


    final Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.MIME_TYPE.key(), "application/plain-text");
    attributes.put("header", "src/test/resources/TestMergeContent/head");
    attributes.put("demarcator", "src/test/resources/TestMergeContent/demarcate");
    attributes.put("footer", "src/test/resources/TestMergeContent/foot");

    runner.enqueue("Hello".getBytes("UTF-8"), attributes);
    runner.enqueue(", ".getBytes("UTF-8"), attributes);
    runner.enqueue("World!".getBytes("UTF-8"), attributes);
    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 testZip() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_BIN_AGE, "1 sec");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_ZIP);

    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);
    try (final InputStream rawIn = new ByteArrayInputStream(runner.getContentAsByteArray(bundle)); final ZipInputStream in = new ZipInputStream(rawIn)) {
        Assert.assertNotNull(in.getNextEntry());
        final byte[] part1 = IOUtils.toByteArray(in);
        Assert.assertTrue(Arrays.equals("Hello".getBytes("UTF-8"), part1));

        in.getNextEntry();
        final byte[] part2 = IOUtils.toByteArray(in);
        Assert.assertTrue(Arrays.equals(", ".getBytes("UTF-8"), part2));

        in.getNextEntry();
        final byte[] part3 = IOUtils.toByteArray(in);
        Assert.assertTrue(Arrays.equals("World!".getBytes("UTF-8"), part3));
    }
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/zip");
}
 

@Test
public void testSmallSplitsThenMerge() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitContent());
    runner.setProperty(SplitContent.KEEP_SEQUENCE, "true");
    runner.setProperty(SplitContent.BYTE_SEQUENCE.getName(), "FFFF");

    runner.enqueue(new byte[]{1, 2, 3, 4, 5, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, 5, 4, 3, 2, 1});
    runner.run();

    runner.assertTransferCount(SplitContent.REL_ORIGINAL, 1);
    runner.getFlowFilesForRelationship(SplitContent.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT, "2");
    runner.assertTransferCount(SplitContent.REL_SPLITS, 2);

    runner.assertQueueEmpty();
    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(SplitContent.REL_SPLITS);
    final MockFlowFile split1 = splits.get(0);
    final MockFlowFile split2 = splits.get(1);

    split1.assertContentEquals(new byte[]{1, 2, 3, 4, 5, (byte) 0xFF, (byte) 0xFF});
    split2.assertContentEquals(new byte[]{(byte) 0xFF, 5, 4, 3, 2, 1});

    final TestRunner mergeRunner = TestRunners.newTestRunner(new MergeContent());
    mergeRunner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_CONCAT);
    mergeRunner.setProperty(MergeContent.MERGE_STRATEGY, MergeContent.MERGE_STRATEGY_DEFRAGMENT);
    mergeRunner.enqueue(splits.toArray(new MockFlowFile[0]));
    mergeRunner.run();

    mergeRunner.assertTransferCount(MergeContent.REL_MERGED, 1);
    mergeRunner.assertTransferCount(MergeContent.REL_ORIGINAL, 2);
    mergeRunner.assertTransferCount(MergeContent.REL_FAILURE, 0);

    final List<MockFlowFile> packed = mergeRunner.getFlowFilesForRelationship(MergeContent.REL_MERGED);
    packed.get(0).assertContentEquals(new byte[]{1, 2, 3, 4, 5, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, 5, 4, 3, 2, 1});
}
 

@Test
public void testRecordResponse() throws IOException, InitializationException {
    SolrClient solrClient = createSolrClient();
    TestRunner runner = createRunnerWithSolrClient(solrClient);

    runner.setProperty(QuerySolr.RETURN_TYPE, QuerySolr.MODE_REC.getValue());
    runner.setProperty(QuerySolr.SOLR_PARAM_FIELD_LIST, "id,created,integer_single");
    runner.setProperty(QuerySolr.SOLR_PARAM_ROWS, "10");

    final String outputSchemaText = new String(Files.readAllBytes(Paths.get("src/test/resources/test-schema.avsc")));

    final JsonRecordSetWriter jsonWriter = new JsonRecordSetWriter();
    runner.addControllerService("writer", jsonWriter);
    runner.setProperty(jsonWriter, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.SCHEMA_TEXT_PROPERTY);
    runner.setProperty(jsonWriter, SchemaAccessUtils.SCHEMA_TEXT, outputSchemaText);
    runner.setProperty(jsonWriter, "Pretty Print JSON", "true");
    runner.setProperty(jsonWriter, "Schema Write Strategy", "full-schema-attribute");
    runner.enableControllerService(jsonWriter);
    runner.setProperty(SolrUtils.RECORD_WRITER, "writer");

    runner.setNonLoopConnection(false);

    runner.run(1);
    runner.assertQueueEmpty();
    runner.assertTransferCount(QuerySolr.RESULTS, 1);

    JsonReader reader = new JsonReader(new InputStreamReader(new ByteArrayInputStream(
            runner.getContentAsByteArray(runner.getFlowFilesForRelationship(QuerySolr.RESULTS).get(0)))));
    reader.beginArray();
    int controlScore = 0;
    while (reader.hasNext()) {
        reader.beginObject();
        while (reader.hasNext()) {
            if (reader.nextName().equals("integer_single")) {
                controlScore += reader.nextInt();
            } else {
                reader.skipValue();
            }
        }
        reader.endObject();
    }
    reader.close();
    solrClient.close();

    assertEquals(controlScore, 45);
}
 

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

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.MIME_TYPE.key(), "application/plain-text");

    attributes.put(CoreAttributes.FILENAME.key(), "AShortFileName");
    runner.enqueue("Hello".getBytes("UTF-8"), attributes);
    attributes.put(CoreAttributes.FILENAME.key(), "ALongerrrFileName");
    runner.enqueue(", ".getBytes("UTF-8"), attributes);
    attributes.put(CoreAttributes.FILENAME.key(), "AReallyLongggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggFileName");
    runner.enqueue("World!".getBytes("UTF-8"), attributes);
    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);
    try (final InputStream rawIn = new ByteArrayInputStream(runner.getContentAsByteArray(bundle)); final TarArchiveInputStream in = new TarArchiveInputStream(rawIn)) {
        ArchiveEntry entry = in.getNextEntry();
        Assert.assertNotNull(entry);
        assertEquals("AShortFileName", entry.getName());
        final byte[] part1 = IOUtils.toByteArray(in);
        Assert.assertTrue(Arrays.equals("Hello".getBytes("UTF-8"), part1));

        entry = in.getNextEntry();
        assertEquals("ALongerrrFileName", entry.getName());
        final byte[] part2 = IOUtils.toByteArray(in);
        Assert.assertTrue(Arrays.equals(", ".getBytes("UTF-8"), part2));

        entry = in.getNextEntry();
        assertEquals("AReallyLongggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggFileName", entry.getName());
        final byte[] part3 = IOUtils.toByteArray(in);
        Assert.assertTrue(Arrays.equals("World!".getBytes("UTF-8"), part3));
    }
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/tar");
}
 

@Test
public void testAvroConcatWithDifferentSchemas() 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);

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

    final GenericRecord record1 = new GenericData.Record(schema1);
    record1.put("name", "Alyssa");
    record1.put("favorite_number", 256);

    final GenericRecord record2 = new GenericData.Record(schema2);
    record2.put("name", "Some Place");

    final GenericRecord record3 = new GenericData.Record(schema1);
    record3.put("name", "John");
    record3.put("favorite_number", 5);
    record3.put("favorite_color", "blue");

    final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema1);
    final ByteArrayOutputStream out1 = serializeAvroRecord(schema1, record1, datumWriter);
    final ByteArrayOutputStream out2 = serializeAvroRecord(schema2, record2, datumWriter);
    final ByteArrayOutputStream out3 = serializeAvroRecord(schema1, record3, datumWriter);

    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, 1);
    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");

    final byte[] data = runner.getContentAsByteArray(bundle);
    final Map<String, GenericRecord> users = getGenericRecordMap(data, schema1, "name");
    Assert.assertEquals(2, users.size());
    Assert.assertTrue(users.containsKey("Alyssa"));
    Assert.assertTrue(users.containsKey("John"));

    final MockFlowFile failure = runner.getFlowFilesForRelationship(MergeContent.REL_FAILURE).get(0);
    final byte[] failureData = runner.getContentAsByteArray(failure);
    final Map<String, GenericRecord> places = getGenericRecordMap(failureData, schema2, "name");
    Assert.assertEquals(1, places.size());
    Assert.assertTrue(places.containsKey("Some Place"));
}