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

The following examples show how to use org.apache.nifi.util.TestRunner#assertQueueEmpty() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: ProtobufEncoderTest.java    From nifi-protobuf-processor with MIT License 6 votes vote down vote up
/**
 * 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"));
}
 
Example 2
Source File: TestSplitContent.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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);
}
 
Example 3
Source File: TestEncodeContent.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up
@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"));
}
 
Example 4
Source File: TestMergeContent.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up
@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");
}
 
Example 5
Source File: TestXMLRecordSetWriter.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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)));
}
 
Example 6
Source File: TestGetSolr.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example 7
Source File: TestBase64EncodeContent.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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"));
}
 
Example 8
Source File: TestSplitContent.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up
@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});
}
 
Example 9
Source File: TestXMLRecordSetWriter.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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)));
}
 
Example 10
Source File: TestMergeContent.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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");
}
 
Example 11
Source File: TestMergeContent.java    From nifi with Apache License 2.0 6 votes vote down vote up
@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");
}
 
Example 12
Source File: TestExtractEmailHeaders.java    From nifi with Apache License 2.0 5 votes vote down vote up
/**
 * 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", "");
}
 
Example 13
Source File: ProtobufDecoderTest.java    From nifi-protobuf-processor with MIT License 5 votes vote down vote up
/**
 * 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);
    }
}
 
Example 14
Source File: TestExtractEmailHeaders.java    From nifi with Apache License 2.0 5 votes vote down vote up
/**
 * 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");
}
 
Example 15
Source File: TestMergeContent.java    From nifi with Apache License 2.0 5 votes vote down vote up
@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");
}
 
Example 16
Source File: TestMergeContent.java    From nifi with Apache License 2.0 5 votes vote down vote up
@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");
}
 
Example 17
Source File: TestSplitContent.java    From nifi with Apache License 2.0 5 votes vote down vote up
@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});
}
 
Example 18
Source File: QuerySolrIT.java    From nifi with Apache License 2.0 4 votes vote down vote up
@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);
}
 
Example 19
Source File: TestMergeContent.java    From nifi with Apache License 2.0 4 votes vote down vote up
@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");
}
 
Example 20
Source File: TestMergeContent.java    From localization_nifi with Apache License 2.0 4 votes vote down vote up
@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"));
}