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

The following examples show how to use org.apache.nifi.util.TestRunner#setThreadCount() . 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: TestExecuteStreamCommand.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up 
@Ignore
@Test
public void testExecuteTouch() throws Exception {
    File testFile = new File("target/test.txt");
    testFile.delete();
    File dummy = new File("src/test/resources/ExecuteCommand/1000bytes.txt");
    final TestRunner controller = TestRunners.newTestRunner(ExecuteStreamCommand.class);
    controller.setValidateExpressionUsage(false);
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.setProperty(ExecuteStreamCommand.WORKING_DIR, "target/xx1");
    controller.setThreadCount(6);
    controller.setProperty(ExecuteStreamCommand.EXECUTION_COMMAND, "c:\\cygwin\\bin\\touch");
    controller.setProperty(ExecuteStreamCommand.EXECUTION_ARGUMENTS, "test.txt");
    controller.assertValid();
    controller.run(6);
    List<MockFlowFile> flowFiles = controller.getFlowFilesForRelationship(ExecuteStreamCommand.OUTPUT_STREAM_RELATIONSHIP);
    assertEquals(5, flowFiles.size());
    assertEquals(0, flowFiles.get(0).getSize());

}
Example 2
Source File: TestUnpackContent.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up 
@Test
public void testThreadSafetyUsingAutoDetect() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new UnpackContent());
    runner.setProperty(UnpackContent.PACKAGING_FORMAT, UnpackContent.PackageFormat.AUTO_DETECT_FORMAT.toString());

    Map<String, String> attrsTar = new HashMap<>(1);
    Map<String, String> attrsFFv3 = new HashMap<>(1);
    attrsTar.put("mime.type", UnpackContent.PackageFormat.TAR_FORMAT.getMimeType());
    attrsFFv3.put("mime.type", UnpackContent.PackageFormat.FLOWFILE_STREAM_FORMAT_V3.getMimeType());

    int numThreads = 50;
    runner.setThreadCount(numThreads);

    for (int i=0; i<numThreads; i++) {
        if (i%2 == 0) {
            runner.enqueue(dataPath.resolve("data.tar"), attrsTar);
        } else {
            runner.enqueue(dataPath.resolve("data.flowfilev3"), attrsFFv3);
        }
    }

    runner.run(numThreads);

    runner.assertTransferCount(UnpackContent.REL_SUCCESS, numThreads*2);
}
Example 3
Source File: TestScanContent.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up 
@Test
public void testBlankLineInDictionaryTextEncoding() throws IOException {
    final String dictionaryWithBlankLine = "Line1\n\nLine3";
    final byte[] dictionaryBytes = dictionaryWithBlankLine.getBytes(ScanContent.UTF8);
    final Path dictionaryPath = Paths.get("target/dictionary");
    Files.write(dictionaryPath, dictionaryBytes, StandardOpenOption.CREATE, StandardOpenOption.WRITE);

    final TestRunner runner = TestRunners.newTestRunner(new ScanContent());
    runner.setThreadCount(1);
    runner.setProperty(ScanContent.DICTIONARY, dictionaryPath.toString());
    runner.setProperty(ScanContent.DICTIONARY_ENCODING, ScanContent.TEXT_ENCODING);

    runner.enqueue(new byte[0]);
    runner.run();
    runner.assertTransferCount(ScanContent.REL_NO_MATCH, 1);

}
Example 4
Source File: TestExecuteStreamCommand.java    From nifi with Apache License 2.0 6 votes vote down vote up 
@Ignore
@Test
public void testExecuteTouch() throws Exception {
    File testFile = new File("target/test.txt");
    testFile.delete();
    File dummy = new File("src/test/resources/ExecuteCommand/1000bytes.txt");
    final TestRunner controller = TestRunners.newTestRunner(ExecuteStreamCommand.class);
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.enqueue(dummy.toPath());
    controller.setProperty(ExecuteStreamCommand.WORKING_DIR, "target/xx1");
    controller.setThreadCount(6);
    controller.setProperty(ExecuteStreamCommand.EXECUTION_COMMAND, "c:\\cygwin\\bin\\touch");
    controller.setProperty(ExecuteStreamCommand.EXECUTION_ARGUMENTS, "test.txt");
    controller.assertValid();
    controller.run(6);
    List<MockFlowFile> flowFiles = controller.getFlowFilesForRelationship(ExecuteStreamCommand.OUTPUT_STREAM_RELATIONSHIP);
    assertEquals(5, flowFiles.size());
    assertEquals(0, flowFiles.get(0).getSize());

}
Example 5
Source File: TestUnpackContent.java    From nifi with Apache License 2.0 6 votes vote down vote up 
@Test
public void testThreadSafetyUsingAutoDetect() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new UnpackContent());
    runner.setProperty(UnpackContent.PACKAGING_FORMAT, UnpackContent.PackageFormat.AUTO_DETECT_FORMAT.toString());

    Map<String, String> attrsTar = new HashMap<>(1);
    Map<String, String> attrsFFv3 = new HashMap<>(1);
    attrsTar.put("mime.type", UnpackContent.PackageFormat.TAR_FORMAT.getMimeType());
    attrsFFv3.put("mime.type", UnpackContent.PackageFormat.FLOWFILE_STREAM_FORMAT_V3.getMimeType());

    int numThreads = 50;
    runner.setThreadCount(numThreads);

    for (int i=0; i<numThreads; i++) {
        if (i%2 == 0) {
            runner.enqueue(dataPath.resolve("data.tar"), attrsTar);
        } else {
            runner.enqueue(dataPath.resolve("data.flowfilev3"), attrsFFv3);
        }
    }

    runner.run(numThreads);

    runner.assertTransferCount(UnpackContent.REL_SUCCESS, numThreads*2);
}
Example 6
Source File: TestScanContent.java    From nifi with Apache License 2.0 6 votes vote down vote up 
@Test
public void testBlankLineInDictionaryTextEncoding() throws IOException {
    final String dictionaryWithBlankLine = "Line1\n\nLine3";
    final byte[] dictionaryBytes = dictionaryWithBlankLine.getBytes(ScanContent.UTF8);
    final Path dictionaryPath = Paths.get("target/dictionary");
    Files.write(dictionaryPath, dictionaryBytes, StandardOpenOption.CREATE, StandardOpenOption.WRITE);

    final TestRunner runner = TestRunners.newTestRunner(new ScanContent());
    runner.setThreadCount(1);
    runner.setProperty(ScanContent.DICTIONARY, dictionaryPath.toString());
    runner.setProperty(ScanContent.DICTIONARY_ENCODING, ScanContent.TEXT_ENCODING);

    runner.enqueue(new byte[0]);
    runner.run();
    runner.assertTransferCount(ScanContent.REL_NO_MATCH, 1);

}
Example 7
Source File: TestIdentifyMimeType.java    From nifi with Apache License 2.0 5 votes vote down vote up 
@Test(expected=AssertionError.class)
public void testOnlyOneCustomMimeConfigSpecified() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new IdentifyMimeType());

    String configFile = "src/test/resources/TestIdentifyMimeType/.customConfig.xml";
    runner.setProperty(IdentifyMimeType.MIME_CONFIG_FILE, configFile);

    String configBody = "foo";
    runner.setProperty(IdentifyMimeType.MIME_CONFIG_BODY, configBody);

    runner.setThreadCount(1);
    runner.run();

}
Example 8
Source File: TestQueryRecord.java    From nifi with Apache License 2.0 5 votes vote down vote up 
@Test
public void testSimple() throws InitializationException, IOException, SQLException {
    final MockRecordParser parser = new MockRecordParser();
    parser.addSchemaField("name", RecordFieldType.STRING);
    parser.addSchemaField("age", RecordFieldType.INT);
    parser.addRecord("Tom", 49);

    final MockRecordWriter writer = new MockRecordWriter("\"name\",\"points\"");

    TestRunner runner = getRunner();
    runner.addControllerService("parser", parser);
    runner.enableControllerService(parser);
    runner.addControllerService("writer", writer);
    runner.enableControllerService(writer);

    runner.setProperty(REL_NAME, "select name, age from FLOWFILE WHERE name <> ''");
    runner.setProperty(QueryRecord.RECORD_READER_FACTORY, "parser");
    runner.setProperty(QueryRecord.RECORD_WRITER_FACTORY, "writer");

    final int numIterations = 1;
    for (int i = 0; i < numIterations; i++) {
        runner.enqueue(new byte[0]);
    }

    runner.setThreadCount(4);
    runner.run(2 * numIterations);

    runner.assertTransferCount(REL_NAME, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(REL_NAME).get(0);
    System.out.println(new String(out.toByteArray()));
    out.assertContentEquals("\"name\",\"points\"\n\"Tom\",\"49\"\n");
}
Example 9
Source File: TestQueryRecord.java    From nifi with Apache License 2.0 5 votes vote down vote up 
@Test
public void testNullable() throws InitializationException, IOException, SQLException {
    final MockRecordParser parser = new MockRecordParser();
    parser.addSchemaField("name", RecordFieldType.STRING, true);
    parser.addSchemaField("age", RecordFieldType.INT, true);
    parser.addRecord("Tom", 49);
    parser.addRecord("Alice", null);
    parser.addRecord(null, 36);

    final MockRecordWriter writer = new MockRecordWriter("\"name\",\"points\"");

    TestRunner runner = getRunner();
    runner.addControllerService("parser", parser);
    runner.enableControllerService(parser);
    runner.addControllerService("writer", writer);
    runner.enableControllerService(writer);

    runner.setProperty(REL_NAME, "select name, age from FLOWFILE");
    runner.setProperty(QueryRecord.RECORD_READER_FACTORY, "parser");
    runner.setProperty(QueryRecord.RECORD_WRITER_FACTORY, "writer");

    final int numIterations = 1;
    for (int i = 0; i < numIterations; i++) {
        runner.enqueue(new byte[0]);
    }

    runner.setThreadCount(4);
    runner.run(2 * numIterations);

    runner.assertTransferCount(REL_NAME, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(REL_NAME).get(0);
    System.out.println(new String(out.toByteArray()));
    out.assertContentEquals("\"name\",\"points\"\n\"Tom\",\"49\"\n\"Alice\",\n,\"36\"\n");
}
Example 10
Source File: TestQueryRecord.java    From nifi with Apache License 2.0 5 votes vote down vote up 
@Test
public void testParseFailure() throws InitializationException, IOException, SQLException {
    final MockRecordParser parser = new MockRecordParser();
    parser.addSchemaField("name", RecordFieldType.STRING);
    parser.addSchemaField("age", RecordFieldType.INT);
    parser.addRecord("Tom", 49);

    final MockRecordWriter writer = new MockRecordWriter("\"name\",\"points\"");

    TestRunner runner = getRunner();
    runner.addControllerService("parser", parser);
    runner.enableControllerService(parser);
    runner.addControllerService("writer", writer);
    runner.enableControllerService(writer);

    runner.setProperty(REL_NAME, "select name, age from FLOWFILE WHERE name <> ''");
    runner.setProperty(QueryRecord.RECORD_READER_FACTORY, "parser");
    runner.setProperty(QueryRecord.RECORD_WRITER_FACTORY, "writer");

    final int numIterations = 1;
    for (int i = 0; i < numIterations; i++) {
        runner.enqueue(new byte[0]);
    }

    runner.setThreadCount(4);
    runner.run(2 * numIterations);

    runner.assertTransferCount(REL_NAME, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(REL_NAME).get(0);
    System.out.println(new String(out.toByteArray()));
    out.assertContentEquals("\"name\",\"points\"\n\"Tom\",\"49\"\n");
}
Example 11
Source File: TestRouteOnAttribute.java    From localization_nifi with Apache License 2.0 4 votes vote down vote up 
@Test
public void testMatchAny() {
    final TestRunner runner = TestRunners.newTestRunner(new RouteOnAttribute());
    runner.setThreadCount(4);
    runner.setProperty(RouteOnAttribute.ROUTE_STRATEGY, RouteOnAttribute.ROUTE_ANY_MATCHES.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(2, match.size());
    assertEquals(2, noMatch.size());

    // Get attributes for both matching FlowFiles
    final Iterator<MockFlowFile> itr = match.iterator();
    final Map<String, String> attrs1 = itr.next().getAttributes();
    final Map<String, String> attrs2 = itr.next().getAttributes();

    // Both matches should map a -> b
    assertEquals("b", attrs1.get("a"));
    assertEquals("b", attrs1.get("a"));

    // One of the flowfiles should map  b -> a, the other b -> b, but we
    // can't know which order they'll come out in, since we're running 4 threads. So either way is acceptable.
    if ("a".equals(attrs1.get("b"))) {
        assertEquals("b", attrs2.get("b"));
    } else {
        assertEquals("b", attrs1.get("b"));
        assertEquals("a", attrs2.get("b"));
    }
    runner.clearTransferState();
}
Example 12
Source File: TestRouteOnAttribute.java    From nifi with Apache License 2.0 4 votes vote down vote up 
@Test
public void testMatchAny() {
    final TestRunner runner = TestRunners.newTestRunner(new RouteOnAttribute());
    runner.setThreadCount(4);
    runner.setProperty(RouteOnAttribute.ROUTE_STRATEGY, RouteOnAttribute.ROUTE_ANY_MATCHES.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(2, match.size());
    assertEquals(2, noMatch.size());

    // Get attributes for both matching FlowFiles
    final Iterator<MockFlowFile> itr = match.iterator();
    final Map<String, String> attrs1 = itr.next().getAttributes();
    final Map<String, String> attrs2 = itr.next().getAttributes();

    // Both matches should map a -> b
    assertEquals("b", attrs1.get("a"));
    assertEquals("b", attrs1.get("a"));

    // One of the flowfiles should map  b -> a, the other b -> b, but we
    // can't know which order they'll come out in, since we're running 4 threads. So either way is acceptable.
    if ("a".equals(attrs1.get("b"))) {
        assertEquals("b", attrs2.get("b"));
    } else {
        assertEquals("b", attrs1.get("b"));
        assertEquals("a", attrs2.get("b"));
    }
    runner.clearTransferState();
}