Examples with ProcessSessionFactory org.apache.nifi.processor.ProcessSessionFactory used on opensource projects

Example 6 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenAndPutSyslog method run.

/**
 * Sends numMessages from PutSyslog to ListenSyslog.
 */
private void run(String protocol, int numMessages, int expectedMessages) throws IOException, InterruptedException {
    // set the same protocol on both processors
    putSyslogRunner.setProperty(PutSyslog.PROTOCOL, protocol);
    listenSyslogRunner.setProperty(ListenSyslog.PROTOCOL, protocol);
    // set a listening port of 0 to get a random available port
    listenSyslogRunner.setProperty(ListenSyslog.PORT, "0");
    // call onScheduled to start ListenSyslog listening
    final ProcessSessionFactory processSessionFactory = listenSyslogRunner.getProcessSessionFactory();
    final ProcessContext context = listenSyslogRunner.getProcessContext();
    listenSyslog.onScheduled(context);
    // get the real port it is listening on and set that in PutSyslog
    final int listeningPort = listenSyslog.getPort();
    putSyslogRunner.setProperty(PutSyslog.PORT, String.valueOf(listeningPort));
    // configure the message properties on PutSyslog
    final String pri = "34";
    final String version = "1";
    final String stamp = "2016-02-05T22:14:15.003Z";
    final String host = "localhost";
    final String body = "some message";
    final String expectedMessage = "<" + pri + ">" + version + " " + stamp + " " + host + " " + body;
    putSyslogRunner.setProperty(PutSyslog.MSG_PRIORITY, pri);
    putSyslogRunner.setProperty(PutSyslog.MSG_VERSION, version);
    putSyslogRunner.setProperty(PutSyslog.MSG_TIMESTAMP, stamp);
    putSyslogRunner.setProperty(PutSyslog.MSG_HOSTNAME, host);
    putSyslogRunner.setProperty(PutSyslog.MSG_BODY, body);
    // send the messages
    for (int i = 0; i < numMessages; i++) {
        putSyslogRunner.enqueue("incoming data".getBytes(Charset.forName("UTF-8")));
    }
    putSyslogRunner.run(numMessages, false);
    // trigger ListenSyslog until we've seen all the messages
    int numTransfered = 0;
    long timeout = System.currentTimeMillis() + 30000;
    while (numTransfered < expectedMessages && System.currentTimeMillis() < timeout) {
        Thread.sleep(10);
        listenSyslog.onTrigger(context, processSessionFactory);
        numTransfered = listenSyslogRunner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).size();
    }
    Assert.assertEquals("Did not process all the messages", expectedMessages, numTransfered);
    if (expectedMessages > 0) {
        // check that one of flow files has the expected content
        MockFlowFile mockFlowFile = listenSyslogRunner.getFlowFilesForRelationship(ListenSyslog.REL_SUCCESS).get(0);
        mockFlowFile.assertContentEquals(expectedMessage);
    }
}

Example 7 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenSyslog method testUDP.

@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();
    }
}

Example 8 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenSyslog method testTCPMultipleConnection.

@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();
    }
}

Example 9 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenSyslog method testTCPSingleConnectionWithNewLines.

@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();
    }
}

Example 10 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class TestConnectWebSocket method testSuccess.

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(ConnectWebSocket.class);
    final ConnectWebSocket processor = (ConnectWebSocket) runner.getProcessor();
    final SharedSessionState sharedSessionState = new SharedSessionState(processor, new AtomicLong(0));
    // Use this custom session factory implementation so that createdSessions can be read from test case,
    // because MockSessionFactory doesn't expose it.
    final Set<MockProcessSession> createdSessions = new HashSet<>();
    final ProcessSessionFactory sessionFactory = () -> {
        final MockProcessSession session = new MockProcessSession(sharedSessionState, processor);
        createdSessions.add(session);
        return session;
    };
    final WebSocketClientService service = mock(WebSocketClientService.class);
    final WebSocketSession webSocketSession = spy(AbstractWebSocketSession.class);
    when(webSocketSession.getSessionId()).thenReturn("ws-session-id");
    when(webSocketSession.getLocalAddress()).thenReturn(new InetSocketAddress("localhost", 12345));
    when(webSocketSession.getRemoteAddress()).thenReturn(new InetSocketAddress("example.com", 80));
    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    when(service.getTargetUri()).thenReturn("ws://example.com/web-socket");
    doAnswer(invocation -> {
        processor.connected(webSocketSession);
        // Two times.
        processor.consume(webSocketSession, textMessageFromServer);
        processor.consume(webSocketSession, textMessageFromServer);
        // Three times.
        final byte[] binaryMessage = textMessageFromServer.getBytes();
        processor.consume(webSocketSession, binaryMessage, 0, binaryMessage.length);
        processor.consume(webSocketSession, binaryMessage, 0, binaryMessage.length);
        processor.consume(webSocketSession, binaryMessage, 0, binaryMessage.length);
        return null;
    }).when(service).connect(endpointId);
    runner.addControllerService(serviceId, service);
    runner.enableControllerService(service);
    runner.setProperty(ConnectWebSocket.PROP_WEBSOCKET_CLIENT_SERVICE, serviceId);
    runner.setProperty(ConnectWebSocket.PROP_WEBSOCKET_CLIENT_ID, endpointId);
    processor.onTrigger(runner.getProcessContext(), sessionFactory);
    final Map<Relationship, List<MockFlowFile>> transferredFlowFiles = getAllTransferredFlowFiles(createdSessions, processor);
    List<MockFlowFile> connectedFlowFiles = transferredFlowFiles.get(AbstractWebSocketGatewayProcessor.REL_CONNECTED);
    assertEquals(1, connectedFlowFiles.size());
    connectedFlowFiles.forEach(ff -> {
        assertFlowFile(webSocketSession, serviceId, endpointId, ff, null);
    });
    List<MockFlowFile> textFlowFiles = transferredFlowFiles.get(AbstractWebSocketGatewayProcessor.REL_MESSAGE_TEXT);
    assertEquals(2, textFlowFiles.size());
    textFlowFiles.forEach(ff -> {
        assertFlowFile(webSocketSession, serviceId, endpointId, ff, WebSocketMessage.Type.TEXT);
    });
    List<MockFlowFile> binaryFlowFiles = transferredFlowFiles.get(AbstractWebSocketGatewayProcessor.REL_MESSAGE_BINARY);
    assertEquals(3, binaryFlowFiles.size());
    binaryFlowFiles.forEach(ff -> {
        assertFlowFile(webSocketSession, serviceId, endpointId, ff, WebSocketMessage.Type.BINARY);
    });
    final List<ProvenanceEventRecord> provenanceEvents = sharedSessionState.getProvenanceEvents();
    assertEquals(6, provenanceEvents.size());
    assertTrue(provenanceEvents.stream().allMatch(event -> ProvenanceEventType.RECEIVE.equals(event.getEventType())));
}