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

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

    final Map<String, String> attrs = new HashMap<>();
    attrs.put("filename", "test1.txt");
    runner.enqueue("test1".getBytes(), attrs);
    runner.run();

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

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    assertEquals(1, events.size());

    final ProvenanceEventRecord event = events.get(0);
    assertEquals(ExtractMediaMetadata.class.getSimpleName(), event.getComponentType());
    assertEquals("media attributes extracted", event.getDetails());
    assertEquals(ProvenanceEventType.ATTRIBUTES_MODIFIED, event.getEventType());
}
 

@Test
public void testSuccessfulDelete() throws Exception {
    Path filePath = new Path("/some/path/to/file.txt");
    when(mockFileSystem.exists(any(Path.class))).thenReturn(true);
    when(mockFileSystem.getUri()).thenReturn(new URI("hdfs://0.example.com:8020"));
    DeleteHDFS deleteHDFS = new TestableDeleteHDFS(kerberosProperties, mockFileSystem);
    TestRunner runner = TestRunners.newTestRunner(deleteHDFS);
    runner.setIncomingConnection(false);
    runner.assertNotValid();
    runner.setProperty(DeleteHDFS.FILE_OR_DIRECTORY, filePath.toString());
    runner.assertValid();
    runner.run();
    // Even if there's no incoming relationship, a FlowFile is created to indicate which path is deleted.
    runner.assertTransferCount(DeleteHDFS.REL_SUCCESS, 1);
    runner.assertTransferCount(DeleteHDFS.REL_FAILURE, 0);

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(1, provenanceEvents.size());
    assertEquals(ProvenanceEventType.REMOTE_INVOCATION, provenanceEvents.get(0).getEventType());
    assertEquals("hdfs://0.example.com:8020/some/path/to/file.txt", provenanceEvents.get(0).getTransitUri());
}
 

@Test
public void testDirectoryUsesValidEL() throws IOException {
    GetHDFS proc = new TestableGetHDFS(kerberosProperties);
    TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(PutHDFS.DIRECTORY, "src/test/resources/${literal('testdata'):substring(0,8)}");
    runner.setProperty(GetHDFS.FILE_FILTER_REGEX, ".*.zip");
    runner.setProperty(GetHDFS.KEEP_SOURCE_FILE, "true");
    runner.setProperty(GetHDFS.COMPRESSION_CODEC, "AUTOMATIC");
    runner.run();

    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(GetHDFS.REL_SUCCESS);
    assertEquals(1, flowFiles.size());

    MockFlowFile flowFile = flowFiles.get(0);
    assertTrue(flowFile.getAttribute(CoreAttributes.FILENAME.key()).equals("13545423550275052.zip"));
    InputStream expected = getClass().getResourceAsStream("/testdata/13545423550275052.zip");
    flowFile.assertContentEquals(expected);
    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(1, provenanceEvents.size());
    final ProvenanceEventRecord receiveEvent = provenanceEvents.get(0);
    assertEquals(ProvenanceEventType.RECEIVE, receiveEvent.getEventType());
    // If it runs with a real HDFS, the protocol will be "hdfs://", but with a local filesystem, just assert the filename.
    assertTrue(receiveEvent.getTransitUri().endsWith("13545423550275052.zip"));
}
 

/**
 * Tests basic ES functionality against a local or test ES cluster
 */
@Test
@Ignore("Comment this out if you want to run against local or test ES")
public void testPutElasticSearchBasic() {
    System.out.println("Starting test " + new Object() {
    }.getClass().getEnclosingMethod().getName());
    final TestRunner runner = TestRunners.newTestRunner(new PutElasticsearchHttpRecord());

    runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "http://127.0.0.1:9200");
    runner.setProperty(PutElasticsearchHttpRecord.INDEX, "doc");
    runner.setProperty(PutElasticsearchHttpRecord.TYPE, "status");
    runner.setProperty(PutElasticsearchHttpRecord.ID_RECORD_PATH, "/id");
    runner.assertValid();

    runner.enqueue(new byte[0], new HashMap<String, String>() {{
        put("doc_id", "28039652140");
    }});

    runner.enqueue(new byte[0]);
    runner.run(1, true, true);
    runner.assertAllFlowFilesTransferred(PutElasticsearchHttpRecord.REL_SUCCESS, 1);
    List<ProvenanceEventRecord> provEvents = runner.getProvenanceEvents();
    assertNotNull(provEvents);
    assertEquals(1, provEvents.size());
    assertEquals(ProvenanceEventType.SEND, provEvents.get(0).getEventType());
}
 

@Test
public void testServiceIsNotFound() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final ControllerService service = spy(ControllerService.class);

    final WebSocketSession webSocketSession = getWebSocketSession();

    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, "different-service-id");
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    runner.enqueue(textMessageFromServer, attributes);

    runner.run();

    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(0, succeededFlowFiles.size());

    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(1, failedFlowFiles.size());
    final MockFlowFile failedFlowFile = failedFlowFiles.iterator().next();
    assertNotNull(failedFlowFile.getAttribute(ATTR_WS_FAILURE_DETAIL));

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());

}
 

@Test
public void testSessionIsNotSpecified() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final WebSocketService service = spy(WebSocketService.class);

    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    runner.enqueue(textMessageFromServer, attributes);

    runner.run();

    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    //assertEquals(0, succeededFlowFiles.size());   //No longer valid test after NIFI-3318 since not specifying sessionid will send to all clients
    assertEquals(1, succeededFlowFiles.size());

    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    //assertEquals(1, failedFlowFiles.size());      //No longer valid test after NIFI-3318
    assertEquals(0, failedFlowFiles.size());

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());

}
 

@Test
public void testSingJsonDocAndExtractedRowId() throws IOException, InitializationException {
    final TestRunner runner = getTestRunner(DEFAULT_TABLE_NAME, DEFAULT_COLUMN_FAMILY, "1");
    final MockHBaseClientService hBaseClient = getHBaseClientService(runner);
    runner.setProperty(PutHBaseJSON.ROW_FIELD_NAME, "rowField");

    final String content = "{ \"rowField\" : \"myRowId\", \"field1\" : \"value1\", \"field2\" : \"value2\" }";
    runner.enqueue(content.getBytes(StandardCharsets.UTF_8));
    runner.run();
    runner.assertAllFlowFilesTransferred(PutHBaseCell.REL_SUCCESS);

    final MockFlowFile outFile = runner.getFlowFilesForRelationship(PutHBaseCell.REL_SUCCESS).get(0);
    outFile.assertContentEquals(content);

    assertNotNull(hBaseClient.getFlowFilePuts());
    assertEquals(1, hBaseClient.getFlowFilePuts().size());

    final List<PutFlowFile> puts = hBaseClient.getFlowFilePuts().get(DEFAULT_TABLE_NAME);
    assertEquals(1, puts.size());

    // should be a put with row id of myRowId, and rowField shouldn't end up in the columns
    final Map<String,byte[]> expectedColumns1 = new HashMap<>();
    expectedColumns1.put("field1", hBaseClient.toBytes("value1"));
    expectedColumns1.put("field2", hBaseClient.toBytes("value2"));
    HBaseTestUtil.verifyPut("myRowId", DEFAULT_COLUMN_FAMILY, expectedColumns1, puts);

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    assertEquals(1, events.size());
    HBaseTestUtil.verifyEvent(runner.getProvenanceEvents(), "hbase://" + DEFAULT_TABLE_NAME + "/myRowId", ProvenanceEventType.SEND);
}
 

@Test
public void testSingleJsonDocAndProvidedRowIdwithNonString() throws IOException, InitializationException {
    final TestRunner runner = getTestRunner(DEFAULT_TABLE_NAME, DEFAULT_COLUMN_FAMILY, "1");
    runner.setProperty(PutHBaseJSON.FIELD_ENCODING_STRATEGY, PutHBaseJSON.BYTES_ENCODING_VALUE);

    final MockHBaseClientService hBaseClient = getHBaseClientService(runner);
    runner.setProperty(PutHBaseJSON.ROW_ID, DEFAULT_ROW);

    final String content = "{ \"field1\" : 1.23456, \"field2\" : 2345235, \"field3\" : false }";
    runner.enqueue(content.getBytes("UTF-8"));
    runner.run();
    runner.assertAllFlowFilesTransferred(PutHBaseCell.REL_SUCCESS);

    final MockFlowFile outFile = runner.getFlowFilesForRelationship(PutHBaseCell.REL_SUCCESS).get(0);
    outFile.assertContentEquals(content);

    assertNotNull(hBaseClient.getFlowFilePuts());
    assertEquals(1, hBaseClient.getFlowFilePuts().size());

    final List<PutFlowFile> puts = hBaseClient.getFlowFilePuts().get(DEFAULT_TABLE_NAME);
    assertEquals(1, puts.size());

    final Map<String,byte[]> expectedColumns = new HashMap<>();
    expectedColumns.put("field1", hBaseClient.toBytes(1.23456d));
    expectedColumns.put("field2", hBaseClient.toBytes(2345235l));
    expectedColumns.put("field3", hBaseClient.toBytes(false));
    HBaseTestUtil.verifyPut(DEFAULT_ROW, DEFAULT_COLUMN_FAMILY, expectedColumns, puts);

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    assertEquals(1, events.size());

    final ProvenanceEventRecord event = events.get(0);
    assertEquals("hbase://" + DEFAULT_TABLE_NAME + "/" + DEFAULT_ROW, event.getTransitUri());
}
 

@Test
public void testSingleJsonDocAndProvidedRowId() throws IOException, InitializationException {
    final TestRunner runner = getTestRunner(DEFAULT_TABLE_NAME, DEFAULT_COLUMN_FAMILY, "1");
    final MockHBaseClientService hBaseClient = getHBaseClientService(runner);
    runner.setProperty(PutHBaseJSON.ROW_ID, DEFAULT_ROW);

    final String content = "{ \"field1\" : \"value1\", \"field2\" : \"value2\" }";
    runner.enqueue(content.getBytes("UTF-8"));
    runner.run();
    runner.assertAllFlowFilesTransferred(PutHBaseCell.REL_SUCCESS);

    final MockFlowFile outFile = runner.getFlowFilesForRelationship(PutHBaseCell.REL_SUCCESS).get(0);
    outFile.assertContentEquals(content);

    assertNotNull(hBaseClient.getFlowFilePuts());
    assertEquals(1, hBaseClient.getFlowFilePuts().size());

    final List<PutFlowFile> puts = hBaseClient.getFlowFilePuts().get(DEFAULT_TABLE_NAME);
    assertEquals(1, puts.size());

    final Map<String,byte[]> expectedColumns = new HashMap<>();
    expectedColumns.put("field1", hBaseClient.toBytes("value1"));
    expectedColumns.put("field2", hBaseClient.toBytes("value2"));
    HBaseTestUtil.verifyPut(DEFAULT_ROW, DEFAULT_COLUMN_FAMILY, expectedColumns, puts);

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    assertEquals(1, events.size());

    final ProvenanceEventRecord event = events.get(0);
    assertEquals("hbase://" + DEFAULT_TABLE_NAME + "/" + DEFAULT_ROW, event.getTransitUri());
}
 

@Test
public void testPutFile() throws IOException {
    PutHDFS proc = new TestablePutHDFS(kerberosProperties, mockFileSystem);
    TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(PutHDFS.DIRECTORY, "target/test-classes");
    runner.setProperty(PutHDFS.CONFLICT_RESOLUTION, "replace");
    try (FileInputStream fis = new FileInputStream("src/test/resources/testdata/randombytes-1")) {
        Map<String, String> attributes = new HashMap<>();
        attributes.put(CoreAttributes.FILENAME.key(), "randombytes-1");
        runner.enqueue(fis, attributes);
        runner.run();
    }

    List<MockFlowFile> failedFlowFiles = runner
            .getFlowFilesForRelationship(new Relationship.Builder().name("failure").build());
    assertTrue(failedFlowFiles.isEmpty());

    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(PutHDFS.REL_SUCCESS);
    assertEquals(1, flowFiles.size());
    MockFlowFile flowFile = flowFiles.get(0);
    assertTrue(mockFileSystem.exists(new Path("target/test-classes/randombytes-1")));
    assertEquals("randombytes-1", flowFile.getAttribute(CoreAttributes.FILENAME.key()));
    assertEquals("target/test-classes", flowFile.getAttribute(PutHDFS.ABSOLUTE_HDFS_PATH_ATTRIBUTE));

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(1, provenanceEvents.size());
    final ProvenanceEventRecord sendEvent = provenanceEvents.get(0);
    assertEquals(ProvenanceEventType.SEND, sendEvent.getEventType());
    // If it runs with a real HDFS, the protocol will be "hdfs://", but with a local filesystem, just assert the filename.
    assertTrue(sendEvent.getTransitUri().endsWith("target/test-classes/randombytes-1"));
}
 

@Test
public void testSingleJsonDocAndProvidedRowIdwithNonString() throws IOException, InitializationException {
    final TestRunner runner = getTestRunner(DEFAULT_TABLE_NAME, DEFAULT_COLUMN_FAMILY, "1");
    runner.setProperty(PutHBaseJSON.FIELD_ENCODING_STRATEGY, PutHBaseJSON.BYTES_ENCODING_VALUE);

    final MockHBaseClientService hBaseClient = getHBaseClientService(runner);
    runner.setProperty(PutHBaseJSON.ROW_ID, DEFAULT_ROW);

    final String content = "{ \"field1\" : 1.23456, \"field2\" : 2345235, \"field3\" : false }";
    runner.enqueue(content.getBytes("UTF-8"));
    runner.run();
    runner.assertAllFlowFilesTransferred(PutHBaseCell.REL_SUCCESS);

    final MockFlowFile outFile = runner.getFlowFilesForRelationship(PutHBaseCell.REL_SUCCESS).get(0);
    outFile.assertContentEquals(content);

    assertNotNull(hBaseClient.getFlowFilePuts());
    assertEquals(1, hBaseClient.getFlowFilePuts().size());

    final List<PutFlowFile> puts = hBaseClient.getFlowFilePuts().get(DEFAULT_TABLE_NAME);
    assertEquals(1, puts.size());

    final Map<String,byte[]> expectedColumns = new HashMap<>();
    expectedColumns.put("field1", hBaseClient.toBytes(1.23456d));
    expectedColumns.put("field2", hBaseClient.toBytes(2345235l));
    expectedColumns.put("field3", hBaseClient.toBytes(false));
    HBaseTestUtil.verifyPut(DEFAULT_ROW, DEFAULT_COLUMN_FAMILY, expectedColumns, puts);

    final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
    assertEquals(1, events.size());

    final ProvenanceEventRecord event = events.get(0);
    assertEquals("hbase://" + DEFAULT_TABLE_NAME + "/" + DEFAULT_ROW, event.getTransitUri());
}
 

@Test
public void testBatching() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.UDP_VALUE.getValue());
    runner.setProperty(ListenSyslog.PORT, "0");
    runner.setProperty(ListenSyslog.MAX_BATCH_SIZE, "25");
    runner.setProperty(ListenSyslog.MESSAGE_DELIMITER, "|");
    runner.setProperty(ListenSyslog.PARSE_MESSAGES, "false");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    // the processor has internal blocking queue with capacity 10 so we have to send
    // less than that since we are sending all messages before the processors ever runs
    final int numMessages = 5;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // write some UDP messages to the port in the background
    final Thread sender = new Thread(new DatagramSender(port, numMessages, 10, VALID_MESSAGE));
    sender.setDaemon(true);
    sender.start();
    sender.join();

    try {
        proc.onTrigger(context, processSessionFactory);
        runner.assertAllFlowFilesTransferred(ListenSyslog.REL_SUCCESS, 1);

        final MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        Assert.assertEquals("0", flowFile.getAttribute(SyslogAttributes.SYSLOG_PORT.key()));
        Assert.assertEquals(ListenSyslog.UDP_VALUE.getValue(), flowFile.getAttribute(SyslogAttributes.SYSLOG_PROTOCOL.key()));
        Assert.assertTrue(!StringUtils.isBlank(flowFile.getAttribute(SyslogAttributes.SYSLOG_SENDER.key())));

        final String content = new String(flowFile.toByteArray(), StandardCharsets.UTF_8);
        final String[] splits = content.split("\\|");
        Assert.assertEquals(numMessages, splits.length);

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(1, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("udp"));
    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}
 

@Test
public void testTCPSingleConnection() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.TCP_VALUE.getValue());
    runner.setProperty(ListenSyslog.PORT, "0");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    // Allow time for the processor to perform its scheduled start
    Thread.sleep(500);

    final int numMessages = 20;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // write some TCP messages to the port in the background
    final Thread sender = new Thread(new SingleConnectionSocketSender(port, numMessages, 10, VALID_MESSAGE_TCP));
    sender.setDaemon(true);
    sender.start();

    // call onTrigger until we read all messages, or 30 seconds passed
    try {
        int nubTransferred = 0;
        long timeout = System.currentTimeMillis() + 30000;

        while (nubTransferred < numMessages && System.currentTimeMillis() < timeout) {
            Thread.sleep(10);
            proc.onTrigger(context, processSessionFactory);
            nubTransferred = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
        }
        Assert.assertEquals("Did not process all the messages", numMessages, nubTransferred);

        MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        checkFlowFile(flowFile, 0, ListenSyslog.TCP_VALUE.getValue());

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(numMessages, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("tcp"));

    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}
 

@Test
public void testUDP() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.UDP_VALUE.getValue());
    runner.setProperty(ListenSyslog.PORT, "0");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    final int numMessages = 20;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // write some UDP messages to the port in the background
    final Thread sender = new Thread(new DatagramSender(port, numMessages, 10, VALID_MESSAGE));
    sender.setDaemon(true);
    sender.start();

    // call onTrigger until we read all datagrams, or 30 seconds passed
    try {
        int numTransferred = 0;
        long timeout = System.currentTimeMillis() + 30000;

        while (numTransferred < numMessages && System.currentTimeMillis() < timeout) {
            Thread.sleep(10);
            proc.onTrigger(context, processSessionFactory);
            numTransferred = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
        }
        Assert.assertEquals("Did not process all the datagrams", numMessages, numTransferred);

        MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        checkFlowFile(flowFile, 0, ListenSyslog.UDP_VALUE.getValue());

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(numMessages, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("udp"));

    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}
 

@Test
public void testUDP() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.UDP_VALUE.getValue());
    runner.setProperty(ListenSyslog.PORT, "0");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    final int numMessages = 20;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // write some UDP messages to the port in the background
    final Thread sender = new Thread(new DatagramSender(port, numMessages, 10, VALID_MESSAGE));
    sender.setDaemon(true);
    sender.start();

    // call onTrigger until we read all datagrams, or 30 seconds passed
    try {
        int numTransferred = 0;
        long timeout = System.currentTimeMillis() + 30000;

        while (numTransferred < numMessages && System.currentTimeMillis() < timeout) {
            Thread.sleep(10);
            proc.onTrigger(context, processSessionFactory);
            numTransferred = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
        }
        Assert.assertEquals("Did not process all the datagrams", numMessages, numTransferred);

        MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        checkFlowFile(flowFile, 0, ListenSyslog.UDP_VALUE.getValue());

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(numMessages, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("udp"));

    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}
 

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final WebSocketService service = spy(WebSocketService.class);

    final WebSocketSession webSocketSession = getWebSocketSession();

    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    doAnswer(invocation -> {
        final SendMessage sendMessage = invocation.getArgumentAt(2, SendMessage.class);
        sendMessage.send(webSocketSession);
        return null;
    }).when(service).sendMessage(anyString(), anyString(), any(SendMessage.class));
    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    runner.setProperty(PutWebSocket.PROP_WS_MESSAGE_TYPE, "${" + ATTR_WS_MESSAGE_TYPE + "}");

    // Enqueue 1st file as Text.
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    attributes.put(ATTR_WS_MESSAGE_TYPE, WebSocketMessage.Type.TEXT.name());
    runner.enqueue(textMessageFromServer, attributes);

    // Enqueue 2nd file as Binary.
    attributes.put(ATTR_WS_MESSAGE_TYPE, WebSocketMessage.Type.BINARY.name());
    runner.enqueue(textMessageFromServer.getBytes(), attributes);

    runner.run(2);

    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(2, succeededFlowFiles.size());
    assertFlowFile(webSocketSession, serviceId, endpointId, succeededFlowFiles.get(0), WebSocketMessage.Type.TEXT);
    assertFlowFile(webSocketSession, serviceId, endpointId, succeededFlowFiles.get(1), WebSocketMessage.Type.BINARY);

    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(0, failedFlowFiles.size());

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(2, provenanceEvents.size());
}
 

@Test
public void testTCPMultipleConnection() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.TCP_VALUE.getValue());
    runner.setProperty(ListenSyslog.MAX_CONNECTIONS, "5");
    runner.setProperty(ListenSyslog.PORT, "0");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    final int numMessages = 20;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // write some TCP messages to the port in the background
    final Thread sender = new Thread(new MultiConnectionSocketSender(port, numMessages, 10, VALID_MESSAGE_TCP));
    sender.setDaemon(true);
    sender.start();

    // call onTrigger until we read all messages, or 30 seconds passed
    try {
        int nubTransferred = 0;
        long timeout = System.currentTimeMillis() + 30000;

        while (nubTransferred < numMessages && System.currentTimeMillis() < timeout) {
            Thread.sleep(10);
            proc.onTrigger(context, processSessionFactory);
            nubTransferred = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
        }
        Assert.assertEquals("Did not process all the messages", numMessages, nubTransferred);

        MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        checkFlowFile(flowFile, 0, ListenSyslog.TCP_VALUE.getValue());

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(numMessages, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("tcp"));

    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}
 

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final WebSocketService service = spy(WebSocketService.class);

    final WebSocketSession webSocketSession = getWebSocketSession();

    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    doAnswer(invocation -> {
        final SendMessage sendMessage = invocation.getArgument(2);
        sendMessage.send(webSocketSession);
        return null;
    }).when(service).sendMessage(anyString(), anyString(), any(SendMessage.class));
    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    runner.setProperty(PutWebSocket.PROP_WS_MESSAGE_TYPE, "${" + ATTR_WS_MESSAGE_TYPE + "}");

    // Enqueue 1st file as Text.
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    attributes.put(ATTR_WS_MESSAGE_TYPE, WebSocketMessage.Type.TEXT.name());
    runner.enqueue(textMessageFromServer, attributes);

    // Enqueue 2nd file as Binary.
    attributes.put(ATTR_WS_MESSAGE_TYPE, WebSocketMessage.Type.BINARY.name());
    runner.enqueue(textMessageFromServer.getBytes(), attributes);

    runner.run(2);

    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(2, succeededFlowFiles.size());
    assertFlowFile(webSocketSession, serviceId, endpointId, succeededFlowFiles.get(0), WebSocketMessage.Type.TEXT);
    assertFlowFile(webSocketSession, serviceId, endpointId, succeededFlowFiles.get(1), WebSocketMessage.Type.BINARY);

    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(0, failedFlowFiles.size());

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(2, provenanceEvents.size());
}
 

@Test
public void testTCPSingleConnection() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.TCP_VALUE.getValue());
    runner.setProperty(ListenSyslog.PORT, "0");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    // Allow time for the processor to perform its scheduled start
    Thread.sleep(500);

    final int numMessages = 20;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // write some TCP messages to the port in the background
    final Thread sender = new Thread(new SingleConnectionSocketSender(port, numMessages, 10, VALID_MESSAGE_TCP));
    sender.setDaemon(true);
    sender.start();

    // call onTrigger until we read all messages, or 30 seconds passed
    try {
        int nubTransferred = 0;
        long timeout = System.currentTimeMillis() + 30000;

        while (nubTransferred < numMessages && System.currentTimeMillis() < timeout) {
            Thread.sleep(10);
            proc.onTrigger(context, processSessionFactory);
            nubTransferred = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
        }
        Assert.assertEquals("Did not process all the messages", numMessages, nubTransferred);

        MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        checkFlowFile(flowFile, 0, ListenSyslog.TCP_VALUE.getValue());

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(numMessages, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("tcp"));

    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}
 

@Test
public void testTCPSingleConnectionWithNewLines() throws IOException, InterruptedException {
    final ListenSyslog proc = new ListenSyslog();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    runner.setProperty(ListenSyslog.PROTOCOL, ListenSyslog.TCP_VALUE.getValue());
    runner.setProperty(ListenSyslog.PORT, "0");

    // schedule to start listening on a random port
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    final ProcessContext context = runner.getProcessContext();
    proc.onScheduled(context);

    final int numMessages = 3;
    final int port = proc.getPort();
    Assert.assertTrue(port > 0);

    // send 3 messages as 1
    final String multipleMessages = VALID_MESSAGE_TCP + "\n" + VALID_MESSAGE_TCP + "\n" + VALID_MESSAGE_TCP + "\n";
    final Thread sender = new Thread(new SingleConnectionSocketSender(port, 1, 10, multipleMessages));
    sender.setDaemon(true);
    sender.start();

    // call onTrigger until we read all messages, or 30 seconds passed
    try {
        int nubTransferred = 0;
        long timeout = System.currentTimeMillis() + 30000;

        while (nubTransferred < numMessages && System.currentTimeMillis() < timeout) {
            Thread.sleep(10);
            proc.onTrigger(context, processSessionFactory);
            nubTransferred = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
        }
        Assert.assertEquals("Did not process all the messages", numMessages, nubTransferred);

        MockFlowFile flowFile = runner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        checkFlowFile(flowFile, 0, ListenSyslog.TCP_VALUE.getValue());

        final List<ProvenanceEventRecord> events = runner.getProvenanceEvents();
        Assert.assertNotNull(events);
        Assert.assertEquals(numMessages, events.size());

        final ProvenanceEventRecord event = events.get(0);
        Assert.assertEquals(ProvenanceEventType.RECEIVE, event.getEventType());
        Assert.assertTrue("transit uri must be set and start with proper protocol", event.getTransitUri().toLowerCase().startsWith("tcp"));

    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
    }
}