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

Example 31 with ProcessSessionFactory

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

the class TestUpdateAttribute method testStateFailures.

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

Example 32 with ProcessSessionFactory

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

the class TestListenWebSocket method testSuccess.

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

Example 33 with ProcessSessionFactory

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

the class ITListenGRPC method testSecureOneWaySSL.

@Test
public void testSecureOneWaySSL() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException, InterruptedException {
    final int randPort = TestGRPCClient.randomPort();
    final Map<String, String> sslProperties = getTruststoreProperties();
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort, sslProperties);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    runner.setProperty(ListenGRPC.PROP_USE_SECURE, "true");
    useSSLContextService(runner, getKeystoreProperties());
    final ProcessContext processContext = runner.getProcessContext();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").setContent(ByteString.copyFrom("content".getBytes())).build();
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.SUCCESS));
        assertThat(reply.getBody(), equalTo("FlowFile successfully received."));
        // known race condition spot: grpc reply vs flowfile transfer
        Thread.sleep(10);
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 1);
        final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ListenGRPC.REL_SUCCESS);
        assertThat(successFiles.size(), equalTo(1));
        final MockFlowFile mockFlowFile = successFiles.get(0);
        assertThat(mockFlowFile.getAttribute("FOO"), equalTo("BAR"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_HOST), equalTo("127.0.0.1"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_USER_DN), equalTo(FlowFileIngestServiceInterceptor.DEFAULT_FOUND_SUBJECT));
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}

Example 34 with ProcessSessionFactory

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

the class ITListenGRPC method testSecureTwoWaySSL.

@Test
public void testSecureTwoWaySSL() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
    final int randPort = TestGRPCClient.randomPort();
    final Map<String, String> sslProperties = getKeystoreProperties();
    sslProperties.putAll(getTruststoreProperties());
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort, sslProperties);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    runner.setProperty(ListenGRPC.PROP_USE_SECURE, "true");
    useSSLContextService(runner, sslProperties);
    final ProcessContext processContext = runner.getProcessContext();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").setContent(ByteString.copyFrom("content".getBytes())).build();
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.SUCCESS));
        assertThat(reply.getBody(), equalTo("FlowFile successfully received."));
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 1);
        final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ListenGRPC.REL_SUCCESS);
        assertThat(successFiles.size(), equalTo(1));
        final MockFlowFile mockFlowFile = successFiles.get(0);
        assertThat(mockFlowFile.getAttribute("FOO"), equalTo("BAR"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_HOST), equalTo("127.0.0.1"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_USER_DN), equalTo(CERT_DN));
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}

Example 35 with ProcessSessionFactory

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

the class ITListenGRPC method testSecureTwoWaySSLFailAuthorizedDNCheck.

@Test(expected = io.grpc.StatusRuntimeException.class)
public void testSecureTwoWaySSLFailAuthorizedDNCheck() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
    final int randPort = TestGRPCClient.randomPort();
    final Map<String, String> sslProperties = getKeystoreProperties();
    sslProperties.putAll(getTruststoreProperties());
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort, sslProperties);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    runner.setProperty(ListenGRPC.PROP_USE_SECURE, "true");
    runner.setProperty(ListenGRPC.PROP_AUTHORIZED_DN_PATTERN, "CN=FAKE.*");
    useSSLContextService(runner, sslProperties);
    final ProcessContext processContext = runner.getProcessContext();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").setContent(ByteString.copyFrom("content".getBytes())).build();
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.SUCCESS));
        assertThat(reply.getBody(), equalTo("FlowFile successfully received."));
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 1);
        final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ListenGRPC.REL_SUCCESS);
        assertThat(successFiles.size(), equalTo(1));
        final MockFlowFile mockFlowFile = successFiles.get(0);
        assertThat(mockFlowFile.getAttribute("FOO"), equalTo("BAR"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_HOST), equalTo("127.0.0.1"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_USER_DN), equalTo(CERT_DN));
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}