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

Example 6 with SharedSessionState

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

the class OldestFirstPrioritizerTest 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 OldestFlowFileFirstPrioritizer prioritizer = new OldestFlowFileFirstPrioritizer();
    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));
}

Example 7 with SharedSessionState

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

the class SNMPUtilsTest method validateUpdateFlowFileAttributes.

/**
 * Test for updating attributes of flow files with {@link PDU}
 */
@Test
public void validateUpdateFlowFileAttributes() {
    GetSNMP processor = new GetSNMP();
    ProcessSession processSession = new MockProcessSession(new SharedSessionState(processor, new AtomicLong()), processor);
    FlowFile sourceFlowFile = processSession.create();
    PDU pdu = new PDU();
    pdu.setErrorIndex(0);
    pdu.setErrorStatus(0);
    pdu.setType(4);
    FlowFile f2 = SNMPUtils.updateFlowFileAttributesWithPduProperties(pdu, sourceFlowFile, processSession);
    assertEquals("0", f2.getAttributes().get(SNMPUtils.SNMP_PROP_PREFIX + "errorIndex"));
    assertEquals("0", f2.getAttributes().get(SNMPUtils.SNMP_PROP_PREFIX + "errorStatus"));
    assertEquals("4", f2.getAttributes().get(SNMPUtils.SNMP_PROP_PREFIX + "type"));
}

Example 8 with SharedSessionState

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

the class TestConsumeAzureEventHub method setupProcessor.

@Before
public void setupProcessor() {
    processor = new ConsumeAzureEventHub();
    final ProcessorInitializationContext initContext = Mockito.mock(ProcessorInitializationContext.class);
    final String componentId = "componentId";
    when(initContext.getIdentifier()).thenReturn(componentId);
    MockComponentLog componentLog = new MockComponentLog(componentId, processor);
    when(initContext.getLogger()).thenReturn(componentLog);
    processor.initialize(initContext);
    final ProcessSessionFactory processSessionFactory = Mockito.mock(ProcessSessionFactory.class);
    processor.setProcessSessionFactory(processSessionFactory);
    processor.setNamespaceName("namespace");
    sharedState = new SharedSessionState(processor, new AtomicLong(0));
    processSession = new MockProcessSession(sharedState, processor);
    when(processSessionFactory.createSession()).thenReturn(processSession);
    eventProcessor = processor.new EventProcessor();
    partitionContext = Mockito.mock(PartitionContext.class);
    when(partitionContext.getEventHubPath()).thenReturn("eventhub-name");
    when(partitionContext.getPartitionId()).thenReturn("partition-id");
    when(partitionContext.getConsumerGroupName()).thenReturn("consumer-group");
}

Example 9 with SharedSessionState

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