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

@Test
public void testRedirectErrorStream() {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setProperty(ExecuteProcess.COMMAND, "cd");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "does-not-exist");
    runner.setProperty(ExecuteProcess.REDIRECT_ERROR_STREAM, "true");

    ProcessContext processContext = runner.getProcessContext();

    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    processor.updateScheduledTrue();
    processor.setupExecutor(processContext);

    processor.onTrigger(processContext, runner.getProcessSessionFactory());

    if (isCommandFailed(runner)) return;

    final List<LogMessage> warnMessages = runner.getLogger().getWarnMessages();
    assertEquals("If redirect error stream is true " +
            "the output should be sent as a content of flow-file.", 0, warnMessages.size());
    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(ExecuteProcess.REL_SUCCESS);
    assertEquals(1, succeeded.size());
}
 

@SuppressWarnings("unchecked")
@Test
public void validateOneWayFromSpring() throws Exception {
    TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
    runner.setProperty(SpringContextProcessor.CTX_CONFIG_PATH, "fromSpringOnly.xml");
    runner.setProperty(SpringContextProcessor.CTX_LIB_PATH, ".");
    runner.assertValid();

    TestProcessor processor = (TestProcessor) runner.getProcessor();
    SpringDataExchanger delegate = processor.getMockedDelegate();
    SpringResponse<Object> r = new SpringResponse<Object>("hello".getBytes(),
            Collections.<String, Object> emptyMap());
    when(delegate.receive(Mockito.anyLong())).thenReturn(r);
    when(delegate.send(Mockito.any(), Mockito.any(Map.class), Mockito.anyLong())).thenReturn(true);

    runner.run(1, false);

    verify(delegate, never()).send(Mockito.any(), Mockito.any(Map.class), Mockito.anyLong());
    verify(delegate, times(1)).receive(0);
    assertTrue(runner.getFlowFilesForRelationship(TestProcessor.REL_SUCCESS).size() == 1);
    runner.shutdown();
}
 

@SuppressWarnings("unchecked")
@Test
public void validateOneWayFromNiFi() throws Exception {
    TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
    runner.setProperty(SpringContextProcessor.CTX_CONFIG_PATH, "toSpringOnly.xml");
    runner.setProperty(SpringContextProcessor.CTX_LIB_PATH, ".");
    runner.setProperty(SpringContextProcessor.RECEIVE_TIMEOUT, "100 millis");
    runner.assertValid();
    runner.enqueue("Hello".getBytes());

    TestProcessor processor = (TestProcessor) runner.getProcessor();
    SpringDataExchanger delegate = processor.getMockedDelegate();
    when(delegate.receive(Mockito.anyLong())).thenReturn(null);
    when(delegate.send(Mockito.any(), Mockito.any(Map.class), Mockito.anyLong())).thenReturn(true);

    runner.run(1, false);
    verify(delegate, times(1)).send(Mockito.any(), Mockito.any(Map.class), Mockito.anyLong());
    verify(delegate, times(1)).receive(100);
    assertTrue(runner.getFlowFilesForRelationship(TestProcessor.REL_SUCCESS).isEmpty());
    runner.shutdown();
}
 

/**
 * 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"));
}
 

@Test
public void testRedirectErrorStreamWithExpressions() {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setProperty(ExecuteProcess.COMMAND, "ls");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "${literal('does-not-exist'):toUpper()}");
    runner.setProperty(ExecuteProcess.REDIRECT_ERROR_STREAM, "true");

    ProcessContext processContext = runner.getProcessContext();

    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    processor.updateScheduledTrue();
    processor.setupExecutor(processContext);

    processor.onTrigger(processContext, runner.getProcessSessionFactory());

    if (isCommandFailed(runner)) return;

    final List<LogMessage> warnMessages = runner.getLogger().getWarnMessages();
    assertEquals("If redirect error stream is true " +
            "the output should be sent as a content of flow-file.", 0, warnMessages.size());
    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(ExecuteProcess.REL_SUCCESS);
    assertEquals(1, succeeded.size());
    assertTrue(new String(succeeded.get(0).toByteArray()).contains("DOES-NOT-EXIST"));
    assertEquals(succeeded.get(0).getAttribute(ExecuteProcess.ATTRIBUTE_COMMAND), "ls");
    assertEquals(succeeded.get(0).getAttribute(ExecuteProcess.ATTRIBUTE_COMMAND_ARGS), "DOES-NOT-EXIST");
}
 

@Test
public void testRedirectErrorStream() {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setProperty(ExecuteProcess.COMMAND, "cd");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "does-not-exist");
    runner.setProperty(ExecuteProcess.REDIRECT_ERROR_STREAM, "true");

    ProcessContext processContext = runner.getProcessContext();

    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    processor.updateScheduledTrue();
    processor.setupExecutor(processContext);

    processor.onTrigger(processContext, runner.getProcessSessionFactory());

    if (isCommandFailed(runner)) return;

    final List<LogMessage> warnMessages = runner.getLogger().getWarnMessages();
    assertEquals("If redirect error stream is true " +
            "the output should be sent as a content of flow-file.", 0, warnMessages.size());
    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(ExecuteProcess.REL_SUCCESS);
    assertEquals(1, succeeded.size());
}
 

@Test
public void validateProcessInterruptOnStop() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setProperty(ExecuteProcess.COMMAND, "ping");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "nifi.apache.org");
    runner.setProperty(ExecuteProcess.BATCH_DURATION, "500 millis");

    runner.run();
    Thread.sleep(500);
    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    try {
        Field executorF = ExecuteProcess.class.getDeclaredField("executor");
        executorF.setAccessible(true);
        ExecutorService executor = (ExecutorService) executorF.get(processor);
        assertTrue(executor.isShutdown());
        assertTrue(executor.isTerminated());

        Field processF = ExecuteProcess.class.getDeclaredField("externalProcess");
        processF.setAccessible(true);
        Process process = (Process) processF.get(processor);
        assertFalse(process.isAlive());
    } catch (Exception e) {
        fail();
    }

}
 

@Test
public void testNotRedirectErrorStream() {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setProperty(ExecuteProcess.COMMAND, "cd");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "does-not-exist");

    ProcessContext processContext = runner.getProcessContext();

    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    processor.updateScheduledTrue();
    processor.setupExecutor(processContext);

    processor.onTrigger(processContext, runner.getProcessSessionFactory());

    if (isCommandFailed(runner)) return;

    // ExecuteProcess doesn't wait for finishing to drain error stream if it's configure NOT to redirect stream.
    // This causes test failure when draining the error stream didn't finish
    // fast enough before the thread of this test case method checks the warn msg count.
    // So, this loop wait for a while until the log msg count becomes expected number, otherwise let it fail.
    final int expectedWarningMessages = 1;
    final int maxRetry = 5;
    for (int i = 0; i < maxRetry
        && (runner.getLogger().getWarnMessages().size() < expectedWarningMessages); i++) {
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e) {
        }
    }
    final List<LogMessage> warnMessages = runner.getLogger().getWarnMessages();
    assertEquals("If redirect error stream is false, " +
            "the output should be logged as a warning so that user can notice on bulletin.", expectedWarningMessages, warnMessages.size());
    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(ExecuteProcess.REL_SUCCESS);
    assertEquals(0, succeeded.size());
}
 

@Test
public void testMap2FlowFileTextMessage() throws Exception {

    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    TextMessage textMessage = new ActiveMQTextMessage();

    String payload = "Hello world!";
    textMessage.setText(payload);

    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());

    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, textMessage, true, pic.getLogger());

    assertEquals("TextMessage content length should equal to FlowFile content size", payload.length(), summary.getLastFlowFile().getSize());

    final byte[] buffer = new byte[payload.length()];
    runner.clearTransferState();

    session.read(summary.getLastFlowFile(), new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer, false);
        }
    });

    String contentString = new String(buffer, "UTF-8");
    assertEquals("", payload, contentString);
}
 

@Test
public void testMap2FlowFileTextMessage() throws Exception {

    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    TextMessage textMessage = new ActiveMQTextMessage();

    String payload = "Hello world!";
    textMessage.setText(payload);

    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());

    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, textMessage, true, pic.getLogger());

    assertEquals("TextMessage content length should equal to FlowFile content size", payload.length(), summary.getLastFlowFile().getSize());

    final byte[] buffer = new byte[payload.length()];
    runner.clearTransferState();

    session.read(summary.getLastFlowFile(), new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer, false);
        }
    });

    String contentString = new String(buffer, "UTF-8");
    assertEquals("", payload, contentString);
}
 

@Test
public void testStateFailures() throws InterruptedException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(AttributeRollingWindow.class);
    MockStateManager mockStateManager = runner.getStateManager();
    final AttributeRollingWindow processor = (AttributeRollingWindow) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();

    runner.setProperty(AttributeRollingWindow.VALUE_TO_TRACK, "${value}");
    runner.setProperty(AttributeRollingWindow.TIME_WINDOW, "3 sec");

    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("value", "1");

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);

    runner.enqueue(new byte[0],attributes);
    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueNotEmpty();

    mockStateManager.setFailOnStateGet(Scope.LOCAL, false);
    mockStateManager.setFailOnStateSet(Scope.LOCAL, true);

    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueEmpty();

    runner.assertAllFlowFilesTransferred(AttributeRollingWindow.REL_FAILED_SET_STATE, 1);
    MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(REL_FAILED_SET_STATE).get(0);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_VALUE_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_COUNT_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_MEAN_KEY);
}
 

@Test
public void validateProcessInterruptOnStop() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setVariable("command", "ping");
    runner.setProperty(ExecuteProcess.COMMAND, "${command}");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "nifi.apache.org");
    runner.setProperty(ExecuteProcess.BATCH_DURATION, "500 millis");

    runner.run();
    Thread.sleep(500);
    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    try {
        Field executorF = ExecuteProcess.class.getDeclaredField("executor");
        executorF.setAccessible(true);
        ExecutorService executor = (ExecutorService) executorF.get(processor);
        assertTrue(executor.isShutdown());
        assertTrue(executor.isTerminated());

        Field processF = ExecuteProcess.class.getDeclaredField("externalProcess");
        processF.setAccessible(true);
        Process process = (Process) processF.get(processor);
        assertFalse(process.isAlive());
    } catch (Exception e) {
        fail();
    }

    final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(ExecuteProcess.REL_SUCCESS);
    if(!flowFiles.isEmpty()) {
        assertTrue(flowFiles.get(0).getAttribute("command").equals("ping"));
    }
}
 

@Test
public void testStateFailures() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new UpdateAttribute());
    final UpdateAttribute processor = (UpdateAttribute) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    MockStateManager mockStateManager = runner.getStateManager();

    runner.setProperty(UpdateAttribute.STORE_STATE, STORE_STATE_LOCALLY);
    runner.setProperty("count", "${getStateValue('count'):plus(1)}");
    runner.setProperty("sum", "${getStateValue('sum'):plus(${pencils})}");
    runner.setProperty(UpdateAttribute.STATEFUL_VARIABLES_INIT_VALUE, "0");

    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes2 = new HashMap<>();
    attributes2.put("pencils", "2");

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);

    runner.enqueue(new byte[0],attributes2);
    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueNotEmpty();

    mockStateManager.setFailOnStateGet(Scope.LOCAL, false);
    mockStateManager.setFailOnStateSet(Scope.LOCAL, true);

    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueEmpty();

    runner.assertAllFlowFilesTransferred(UpdateAttribute.REL_FAILED_SET_STATE, 1);
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeEquals("count", "1");
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeEquals("sum", "2");
}
 

@SuppressWarnings("unchecked")
@Test
public void validateBiDirectional() throws Exception {
    TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
    runner.setProperty(SpringContextProcessor.CTX_CONFIG_PATH, "requestReply.xml");
    runner.setProperty(SpringContextProcessor.CTX_LIB_PATH, ".");
    runner.setProperty(SpringContextProcessor.RECEIVE_TIMEOUT, "100 millis");
    runner.assertValid();
    runner.enqueue("Hello".getBytes());

    TestProcessor processor = (TestProcessor) runner.getProcessor();
    SpringDataExchanger delegate = processor.getMockedDelegate();
    Map<String, Object> headers = new HashMap<>();
    headers.put("foo", "foo");
    headers.put("bar", new Object());
    SpringResponse<Object> r = new SpringResponse<Object>("hello".getBytes(), headers);
    when(delegate.receive(Mockito.anyLong())).thenReturn(r);
    when(delegate.send(Mockito.any(), Mockito.any(Map.class), Mockito.anyLong())).thenReturn(true);

    runner.run(1, false);
    verify(delegate, times(1)).send(Mockito.any(), Mockito.any(Map.class), Mockito.anyLong());
    verify(delegate, times(1)).receive(100);
    List<MockFlowFile> ffList = runner.getFlowFilesForRelationship(TestProcessor.REL_SUCCESS);
    assertTrue(ffList.size() == 1);
    assertEquals("foo", ffList.get(0).getAttribute("foo"));
    assertNull(ffList.get(0).getAttribute("bar"));
    runner.shutdown();
}
 

@Test
public void testStateFailures() throws InterruptedException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(AttributeRollingWindow.class);
    MockStateManager mockStateManager = runner.getStateManager();
    final AttributeRollingWindow processor = (AttributeRollingWindow) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();

    runner.setProperty(AttributeRollingWindow.VALUE_TO_TRACK, "${value}");
    runner.setProperty(AttributeRollingWindow.TIME_WINDOW, "3 sec");

    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("value", "1");

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);

    runner.enqueue(new byte[0],attributes);
    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueNotEmpty();

    mockStateManager.setFailOnStateGet(Scope.LOCAL, false);
    mockStateManager.setFailOnStateSet(Scope.LOCAL, true);

    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueEmpty();

    runner.assertAllFlowFilesTransferred(AttributeRollingWindow.REL_FAILED_SET_STATE, 1);
    MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(REL_FAILED_SET_STATE).get(0);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_VALUE_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_COUNT_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_MEAN_KEY);
}
 

/**
 * Test BytesMessage to FlowFile conversion
 *
 * @throws java.lang.Exception ex
 */
@Test
public void testMap2FlowFileBytesMessage() throws Exception {

    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    BytesMessage bytesMessage = new ActiveMQBytesMessage();

    String sourceString = "Apache NiFi is an easy to use, powerful, and reliable system to process and distribute data.!";
    byte[] payload = sourceString.getBytes("UTF-8");
    bytesMessage.writeBytes(payload);
    bytesMessage.reset();

    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());

    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, bytesMessage, true, pic.getLogger());

    assertEquals("BytesMessage content length should equal to FlowFile content size", payload.length, summary.getLastFlowFile().getSize());

    final byte[] buffer = new byte[payload.length];
    runner.clearTransferState();

    session.read(summary.getLastFlowFile(), new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer, false);
        }
    });

    String contentString = new String(buffer, "UTF-8");
    assertEquals("", sourceString, contentString);
}
 

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(ListenWebSocket.class);
    final ListenWebSocket processor = (ListenWebSocket)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 WebSocketServerService service = mock(WebSocketServerService.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 = "/test";
    final String textMessageReceived = "message from server.";
    final AtomicReference<Boolean> registered = new AtomicReference<>(false);
    when(service.getIdentifier()).thenReturn(serviceId);
    doAnswer(invocation -> {
        registered.set(true);
        processor.connected(webSocketSession);
        // Two times.
        processor.consume(webSocketSession, textMessageReceived);
        processor.consume(webSocketSession, textMessageReceived);
        // Three times.
        final byte[] binaryMessage = textMessageReceived.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).registerProcessor(endpointId, processor);
    doAnswer(invocation -> registered.get())
            .when(service).isProcessorRegistered(eq(endpointId), eq(processor));
    doAnswer(invocation -> {
        registered.set(false);
        return null;
    }).when(service).deregisterProcessor(eq(endpointId), eq(processor));

    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    runner.setProperty(ListenWebSocket.PROP_WEBSOCKET_SERVER_SERVICE, serviceId);
    runner.setProperty(ListenWebSocket.PROP_SERVER_URL_PATH, endpointId);

    processor.onTrigger(runner.getProcessContext(), sessionFactory);

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

    runner.clearTransferState();
    runner.clearProvenanceEvents();
    createdSessions.clear();
    assertEquals(0, createdSessions.size());

    // Simulate that the processor has started, and it get's triggered again
    processor.onTrigger(runner.getProcessContext(), sessionFactory);
    assertEquals("No session should be created", 0, createdSessions.size());

    // Simulate that the processor is stopped.
    processor.onStopped(runner.getProcessContext());
    assertEquals("No session should be created", 0, createdSessions.size());

    // Simulate that the processor is restarted.
    // And the mock service will emit consume msg events.
    processor.onTrigger(runner.getProcessContext(), sessionFactory);
    assertEquals("Processor should register it with the service again", 6, createdSessions.size());
}
 

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(ListenWebSocket.class);
    final ListenWebSocket processor = (ListenWebSocket)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 WebSocketServerService service = mock(WebSocketServerService.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 = "/test";
    final String textMessageReceived = "message from server.";
    final AtomicReference<Boolean> registered = new AtomicReference<>(false);
    when(service.getIdentifier()).thenReturn(serviceId);
    doAnswer(invocation -> {
        registered.set(true);
        processor.connected(webSocketSession);
        // Two times.
        processor.consume(webSocketSession, textMessageReceived);
        processor.consume(webSocketSession, textMessageReceived);
        // Three times.
        final byte[] binaryMessage = textMessageReceived.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).registerProcessor(endpointId, processor);
    doAnswer(invocation -> registered.get())
            .when(service).isProcessorRegistered(eq(endpointId), eq(processor));
    doAnswer(invocation -> {
        registered.set(false);
        return null;
    }).when(service).deregisterProcessor(eq(endpointId), eq(processor));

    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    runner.setProperty(ListenWebSocket.PROP_WEBSOCKET_SERVER_SERVICE, serviceId);
    runner.setProperty(ListenWebSocket.PROP_SERVER_URL_PATH, endpointId);

    processor.onTrigger(runner.getProcessContext(), sessionFactory);

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

    runner.clearTransferState();
    runner.clearProvenanceEvents();
    createdSessions.clear();
    assertEquals(0, createdSessions.size());

    // Simulate that the processor has started, and it get's triggered again
    processor.onTrigger(runner.getProcessContext(), sessionFactory);
    assertEquals("No session should be created", 0, createdSessions.size());

    // Simulate that the processor is stopped.
    processor.onStopped(runner.getProcessContext());
    assertEquals("No session should be created", 0, createdSessions.size());

    // Simulate that the processor is restarted.
    // And the mock service will emit consume msg events.
    processor.onTrigger(runner.getProcessContext(), sessionFactory);
    assertEquals("Processor should register it with the service again", 6, createdSessions.size());
}
 

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

@Test
public void testStateFailuresWithRulesUsingClone() throws Exception {
    final Criteria criteria = getCriteria();
    criteria.setFlowFilePolicy(FlowFilePolicy.USE_CLONE);
    addRule(criteria, "rule", Collections.singletonList(
            // conditions
            "${getStateValue('maxValue'):lt(${value})}"), getMap(
            // actions
            "maxValue", "${value}"));
    addRule(criteria, "rule2", Collections.singletonList(
            // conditions
            "${getStateValue('maxValue2'):lt(${value})}"), getMap(
            // actions
            "maxValue2", "${value}"));

    TestRunner runner = TestRunners.newTestRunner(new UpdateAttribute());
    final UpdateAttribute processor = (UpdateAttribute) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    MockStateManager mockStateManager = runner.getStateManager();

    runner.setProperty(UpdateAttribute.STORE_STATE, STORE_STATE_LOCALLY);
    runner.setProperty(UpdateAttribute.STATEFUL_VARIABLES_INIT_VALUE, "0");
    runner.setAnnotationData(serialize(criteria));


    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("value", "1");
    runner.enqueue(new byte[0], attributes);

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);
    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueNotEmpty();
    mockStateManager.setFailOnStateGet(Scope.LOCAL, false);
    mockStateManager.setFailOnStateSet(Scope.LOCAL, true);

    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());

    runner.assertQueueEmpty();

    runner.assertAllFlowFilesTransferred(UpdateAttribute.REL_FAILED_SET_STATE, 1);
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeEquals("maxValue", "1");
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeNotExists("maxValue2");
}