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

Example 36 with ProcessContext

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

the class TestJmsConsumer method testMap2FlowFileMapMessage.

/**
 * Test MapMessage to FlowFile conversion
 *
 * @throws java.lang.Exception ex
 */
@Test
public void testMap2FlowFileMapMessage() throws Exception {
    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    MapMessage mapMessage = createMapMessage();
    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());
    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, mapMessage, true, pic.getLogger());
    assertEquals("MapMessage should not create FlowFile content", 0, summary.getBytesReceived());
    Map<String, String> attributes = summary.getLastFlowFile().getAttributes();
    assertEquals("", "Arnold", attributes.get(JmsConsumer.MAP_MESSAGE_PREFIX + "name"));
    assertEquals("", "97", attributes.get(JmsConsumer.MAP_MESSAGE_PREFIX + "age"));
    assertEquals("", "89686.564", attributes.get(JmsConsumer.MAP_MESSAGE_PREFIX + "xyz"));
    assertEquals("", "true", attributes.get(JmsConsumer.MAP_MESSAGE_PREFIX + "good"));
}

Example 37 with ProcessContext

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

the class TestListenSyslog method testParsingError.

@Test
public void testParsingError() throws IOException {
    final FailParseProcessor proc = new FailParseProcessor();
    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);
    try {
        final int port = proc.getPort();
        final DatagramSender sender = new DatagramSender(port, 1, 1, INVALID_MESSAGE);
        sender.run();
        // should keep re-processing event1 from the error queue
        proc.onTrigger(context, processSessionFactory);
        runner.assertTransferCount(ListenSyslog.REL_INVALID, 1);
        runner.assertTransferCount(ListenSyslog.REL_SUCCESS, 0);
    } finally {
        proc.onUnscheduled();
    }
}

Example 38 with ProcessContext

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

the class TestListenSyslog method testBatching.

@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.PORT.key()));
        Assert.assertEquals(ListenSyslog.UDP_VALUE.getValue(), flowFile.getAttribute(SyslogAttributes.PROTOCOL.key()));
        Assert.assertTrue(!StringUtils.isBlank(flowFile.getAttribute(SyslogAttributes.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();
    }
}

Example 39 with ProcessContext

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

the class TestListenTCP method runTCP.

protected void runTCP(final List<String> messages, final int expectedTransferred, final SSLContext sslContext) throws IOException, InterruptedException {
    Socket socket = null;
    try {
        // schedule to start listening on a random port
        final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
        final ProcessContext context = runner.getProcessContext();
        proc.onScheduled(context);
        // create a client connection to the port the dispatcher is listening on
        final int realPort = proc.getDispatcherPort();
        // create either a regular socket or ssl socket based on context being passed in
        if (sslContext != null) {
            socket = sslContext.getSocketFactory().createSocket("localhost", realPort);
        } else {
            socket = new Socket("localhost", realPort);
        }
        Thread.sleep(100);
        // send the frames to the port the processors is listening on
        for (final String message : messages) {
            socket.getOutputStream().write(message.getBytes(StandardCharsets.UTF_8));
            Thread.sleep(1);
        }
        socket.getOutputStream().flush();
        long responseTimeout = 10000;
        // this first loop waits until the internal queue of the processor has the expected
        // number of messages ready before proceeding, we want to guarantee they are all there
        // before onTrigger gets a chance to run
        long startTimeQueueSizeCheck = System.currentTimeMillis();
        while (proc.getQueueSize() < messages.size() && (System.currentTimeMillis() - startTimeQueueSizeCheck < responseTimeout)) {
            Thread.sleep(100);
        }
        // want to fail here if the queue size isn't what we expect
        Assert.assertEquals(messages.size(), proc.getQueueSize());
        // call onTrigger until we processed all the frames, or a certain amount of time passes
        int numTransferred = 0;
        long startTime = System.currentTimeMillis();
        while (numTransferred < expectedTransferred && (System.currentTimeMillis() - startTime < responseTimeout)) {
            proc.onTrigger(context, processSessionFactory);
            numTransferred = runner.getFlowFilesForRelationship(ListenTCP.REL_SUCCESS).size();
            Thread.sleep(100);
        }
        // should have transferred the expected events
        runner.assertTransferCount(ListenTCP.REL_SUCCESS, expectedTransferred);
    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
        IOUtils.closeQuietly(socket);
    }
}

Example 40 with ProcessContext

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

the class TestListenUDP method run.

protected void run(final DatagramSocket socket, final List<String> messages, final int expectedQueueSize, final int expectedTransferred) throws IOException, InterruptedException {
    try {
        // schedule to start listening on a random port
        final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
        final ProcessContext context = runner.getProcessContext();
        proc.onScheduled(context);
        Thread.sleep(100);
        // get the real port the dispatcher is listening on
        final int destPort = proc.getDispatcherPort();
        final InetSocketAddress destination = new InetSocketAddress("localhost", destPort);
        // send the messages to the port the processors is listening on
        for (final String message : messages) {
            final byte[] buffer = message.getBytes(StandardCharsets.UTF_8);
            final DatagramPacket packet = new DatagramPacket(buffer, buffer.length, destination);
            socket.send(packet);
            Thread.sleep(10);
        }
        long responseTimeout = 10000;
        // this first loop waits until the internal queue of the processor has the expected
        // number of messages ready before proceeding, we want to guarantee they are all there
        // before onTrigger gets a chance to run
        long startTimeQueueSizeCheck = System.currentTimeMillis();
        while (proc.getQueueSize() < expectedQueueSize && (System.currentTimeMillis() - startTimeQueueSizeCheck < responseTimeout)) {
            Thread.sleep(100);
        }
        // want to fail here if the queue size isn't what we expect
        Assert.assertEquals(expectedQueueSize, proc.getQueueSize());
        // call onTrigger until we processed all the messages, or a certain amount of time passes
        int numTransferred = 0;
        long startTime = System.currentTimeMillis();
        while (numTransferred < expectedTransferred && (System.currentTimeMillis() - startTime < responseTimeout)) {
            proc.onTrigger(context, processSessionFactory);
            numTransferred = runner.getFlowFilesForRelationship(ListenUDP.REL_SUCCESS).size();
            Thread.sleep(100);
        }
        // should have transferred the expected events
        runner.assertTransferCount(ListenUDP.REL_SUCCESS, expectedTransferred);
    } finally {
        // unschedule to close connections
        proc.onUnscheduled();
        IOUtils.closeQuietly(socket);
    }
}