Examples with SharedSessionState org.apache.nifi.util.SharedSessionState used on opensource projects

Example 1 with SharedSessionState

use of org.apache.nifi.util.SharedSessionState in project nifi by apache.

the class TestStandardRemoteGroupPort method setupMockProcessSession.

private void setupMockProcessSession() {
    // Construct a RemoteGroupPort as a processor to use NiFi mock library.
    final Processor remoteGroupPort = mock(Processor.class);
    final Set<Relationship> relationships = new HashSet<>();
    relationships.add(Relationship.ANONYMOUS);
    when(remoteGroupPort.getRelationships()).thenReturn(relationships);
    when(remoteGroupPort.getIdentifier()).thenReturn("remote-group-port-id");
    sessionState = new SharedSessionState(remoteGroupPort, new AtomicLong(0));
    processSession = new MockProcessSession(sessionState, remoteGroupPort);
    processContext = new MockProcessContext(remoteGroupPort);
}

Example 2 with SharedSessionState

use of org.apache.nifi.util.SharedSessionState in project nifi by apache.

the class PriorityAttributePrioritizerTest method testPrioritizer.

@Test
public void testPrioritizer() throws InstantiationException, IllegalAccessException {
    final Processor processor = new SimpleProcessor();
    final AtomicLong idGenerator = new AtomicLong(0L);
    final MockProcessSession session = new MockProcessSession(new SharedSessionState(processor, idGenerator), Mockito.mock(Processor.class));
    final MockFlowFile ffNoPriority = session.create();
    final MockFlowFile ffPri1 = session.create();
    ffPri1.putAttributes(attrsPri1);
    final MockFlowFile ffPri2 = session.create();
    ffPri2.putAttributes(attrsPri2);
    final MockFlowFile ffPrin1 = session.create();
    ffPrin1.putAttributes(attrsPrin1);
    final MockFlowFile ffPriA = session.create();
    ffPriA.putAttributes(attrsPriA);
    final MockFlowFile ffPriB = session.create();
    ffPriB.putAttributes(attrsPriB);
    final MockFlowFile ffPriLP = session.create();
    ffPriLP.putAttributes(attrsPriLP);
    final MockFlowFile ffPriLN = session.create();
    ffPriLN.putAttributes(attrsPriLN);
    final PriorityAttributePrioritizer prioritizer = new PriorityAttributePrioritizer();
    assertEquals(0, prioritizer.compare(null, null));
    assertEquals(-1, prioritizer.compare(ffNoPriority, null));
    assertEquals(1, prioritizer.compare(null, ffNoPriority));
    assertEquals(0, prioritizer.compare(ffNoPriority, ffNoPriority));
    assertEquals(-1, prioritizer.compare(ffPri1, ffNoPriority));
    assertEquals(1, prioritizer.compare(ffNoPriority, ffPri1));
    assertEquals(0, prioritizer.compare(ffPri1, ffPri1));
    assertEquals(-1, prioritizer.compare(ffPri1, ffPri2));
    assertEquals(1, prioritizer.compare(ffPri2, ffPri1));
    assertEquals(-1, prioritizer.compare(ffPrin1, ffPri1));
    assertEquals(1, prioritizer.compare(ffPri1, ffPrin1));
    assertEquals(-1, prioritizer.compare(ffPri1, ffPriA));
    assertEquals(1, prioritizer.compare(ffPriA, ffPri1));
    assertEquals(0, prioritizer.compare(ffPriA, ffPriA));
    assertEquals(-1, prioritizer.compare(ffPriA, ffPriB));
    assertEquals(1, prioritizer.compare(ffPriB, ffPriA));
    assertEquals(1, prioritizer.compare(ffPriLP, ffPri1));
    assertEquals(-1, prioritizer.compare(ffPri1, ffPriLP));
    assertEquals(-1, prioritizer.compare(ffPriLN, ffPri1));
    assertEquals(1, prioritizer.compare(ffPri1, ffPriLN));
}

Example 3 with SharedSessionState

use of org.apache.nifi.util.SharedSessionState 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())));
}

Example 4 with SharedSessionState

use of org.apache.nifi.util.SharedSessionState in project nifi by apache.

the class TestHttpFlowFileServerProtocol method setupMockProcessSession.

private void setupMockProcessSession() {
    // Construct a RootGroupPort as a processor to use NiFi mock library.
    final Processor rootGroupPort = mock(Processor.class);
    final Set<Relationship> relationships = new HashSet<>();
    relationships.add(Relationship.ANONYMOUS);
    when(rootGroupPort.getRelationships()).thenReturn(relationships);
    when(rootGroupPort.getIdentifier()).thenReturn("root-group-port-id");
    sessionState = new SharedSessionState(rootGroupPort, new AtomicLong(0));
    processSession = new MockProcessSession(sessionState, rootGroupPort);
    processContext = new MockProcessContext(rootGroupPort);
}

Example 5 with SharedSessionState

use of org.apache.nifi.util.SharedSessionState in project nifi by apache.

the class NewestFirstPrioritizerTest method testPrioritizer.

@Test
public void testPrioritizer() throws InstantiationException, IllegalAccessException {
    final Processor processor = new SimpleProcessor();
    final AtomicLong idGenerator = new AtomicLong(0L);
    final MockProcessSession session = new MockProcessSession(new SharedSessionState(processor, idGenerator), Mockito.mock(Processor.class));
    final MockFlowFile flowFile1 = session.create();
    try {
        // guarantee the FlowFile entryDate for flowFile2 is different than flowFile1
        Thread.sleep(2);
    } catch (final InterruptedException e) {
    }
    final MockFlowFile flowFile2 = session.create();
    final NewestFlowFileFirstPrioritizer prioritizer = new NewestFlowFileFirstPrioritizer();
    Assert.assertEquals(0, prioritizer.compare(null, null));
    Assert.assertEquals(-1, prioritizer.compare(flowFile1, null));
    Assert.assertEquals(1, prioritizer.compare(null, flowFile1));
    Assert.assertEquals(0, prioritizer.compare(flowFile1, flowFile1));
    Assert.assertEquals(0, prioritizer.compare(flowFile2, flowFile2));
    Assert.assertEquals(1, prioritizer.compare(flowFile1, flowFile2));
    Assert.assertEquals(-1, prioritizer.compare(flowFile2, flowFile1));
}