Search in sources :

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));
}
Also used : MockFlowFile(org.apache.nifi.util.MockFlowFile) AtomicLong(java.util.concurrent.atomic.AtomicLong) SharedSessionState(org.apache.nifi.util.SharedSessionState) Processor(org.apache.nifi.processor.Processor) AbstractProcessor(org.apache.nifi.processor.AbstractProcessor) MockProcessSession(org.apache.nifi.util.MockProcessSession) Test(org.junit.Test)

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"));
}
Also used : MockProcessSession(org.apache.nifi.util.MockProcessSession) ProcessSession(org.apache.nifi.processor.ProcessSession) PDU(org.snmp4j.PDU) FlowFile(org.apache.nifi.flowfile.FlowFile) SharedSessionState(org.apache.nifi.util.SharedSessionState) AtomicLong(java.util.concurrent.atomic.AtomicLong) MockProcessSession(org.apache.nifi.util.MockProcessSession) Test(org.junit.Test)

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");
}
Also used : PartitionContext(com.microsoft.azure.eventprocessorhost.PartitionContext) SharedSessionState(org.apache.nifi.util.SharedSessionState) AtomicLong(java.util.concurrent.atomic.AtomicLong) MockComponentLog(org.apache.nifi.util.MockComponentLog) ProcessSessionFactory(org.apache.nifi.processor.ProcessSessionFactory) MockProcessSession(org.apache.nifi.util.MockProcessSession) ProcessorInitializationContext(org.apache.nifi.processor.ProcessorInitializationContext) Before(org.junit.Before)

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());
}
Also used : WebSocketServerService(org.apache.nifi.websocket.WebSocketServerService) ATTR_WS_LOCAL_ADDRESS(org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_LOCAL_ADDRESS) HashMap(java.util.HashMap) Mockito.spy(org.mockito.Mockito.spy) AtomicReference(java.util.concurrent.atomic.AtomicReference) ATTR_WS_SESSION_ID(org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_SESSION_ID) ArrayList(java.util.ArrayList) HashSet(java.util.HashSet) Relationship(org.apache.nifi.processor.Relationship) TestRunner(org.apache.nifi.util.TestRunner) ProvenanceEventRecord(org.apache.nifi.provenance.ProvenanceEventRecord) Matchers.eq(org.mockito.Matchers.eq) Map(java.util.Map) Mockito.doAnswer(org.mockito.Mockito.doAnswer) Assert.fail(org.junit.Assert.fail) WebSocketSession(org.apache.nifi.websocket.WebSocketSession) MockProcessSession(org.apache.nifi.util.MockProcessSession) ATTR_WS_ENDPOINT_ID(org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_ENDPOINT_ID) ProvenanceEventType(org.apache.nifi.provenance.ProvenanceEventType) Collection(java.util.Collection) WebSocketMessage(org.apache.nifi.websocket.WebSocketMessage) ATTR_WS_REMOTE_ADDRESS(org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_REMOTE_ADDRESS) Set(java.util.Set) Assert.assertTrue(org.junit.Assert.assertTrue) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) ProcessSessionFactory(org.apache.nifi.processor.ProcessSessionFactory) InetSocketAddress(java.net.InetSocketAddress) Processor(org.apache.nifi.processor.Processor) AtomicLong(java.util.concurrent.atomic.AtomicLong) List(java.util.List) SharedSessionState(org.apache.nifi.util.SharedSessionState) AbstractWebSocketSession(org.apache.nifi.websocket.AbstractWebSocketSession) ATTR_WS_CS_ID(org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_CS_ID) ATTR_WS_MESSAGE_TYPE(org.apache.nifi.processors.websocket.WebSocketProcessorAttributes.ATTR_WS_MESSAGE_TYPE) TestRunners(org.apache.nifi.util.TestRunners) Assert.assertEquals(org.junit.Assert.assertEquals) MockFlowFile(org.apache.nifi.util.MockFlowFile) Mockito.mock(org.mockito.Mockito.mock) TestRunner(org.apache.nifi.util.TestRunner) InetSocketAddress(java.net.InetSocketAddress) AtomicReference(java.util.concurrent.atomic.AtomicReference) WebSocketSession(org.apache.nifi.websocket.WebSocketSession) AbstractWebSocketSession(org.apache.nifi.websocket.AbstractWebSocketSession) MockFlowFile(org.apache.nifi.util.MockFlowFile) SharedSessionState(org.apache.nifi.util.SharedSessionState) AtomicLong(java.util.concurrent.atomic.AtomicLong) WebSocketServerService(org.apache.nifi.websocket.WebSocketServerService) Relationship(org.apache.nifi.processor.Relationship) ProvenanceEventRecord(org.apache.nifi.provenance.ProvenanceEventRecord) ProcessSessionFactory(org.apache.nifi.processor.ProcessSessionFactory) ArrayList(java.util.ArrayList) List(java.util.List) MockProcessSession(org.apache.nifi.util.MockProcessSession) HashSet(java.util.HashSet) Test(org.junit.Test)

Aggregations

AtomicLong (java.util.concurrent.atomic.AtomicLong)9 MockProcessSession (org.apache.nifi.util.MockProcessSession)9 SharedSessionState (org.apache.nifi.util.SharedSessionState)9 Processor (org.apache.nifi.processor.Processor)6 Test (org.junit.Test)6 MockFlowFile (org.apache.nifi.util.MockFlowFile)5 HashSet (java.util.HashSet)4 Relationship (org.apache.nifi.processor.Relationship)4 AbstractProcessor (org.apache.nifi.processor.AbstractProcessor)3 ProcessSessionFactory (org.apache.nifi.processor.ProcessSessionFactory)3 InetSocketAddress (java.net.InetSocketAddress)2 List (java.util.List)2 Map (java.util.Map)2 Set (java.util.Set)2 ProvenanceEventRecord (org.apache.nifi.provenance.ProvenanceEventRecord)2 ProvenanceEventType (org.apache.nifi.provenance.ProvenanceEventType)2 MockProcessContext (org.apache.nifi.util.MockProcessContext)2 TestRunner (org.apache.nifi.util.TestRunner)2 TestRunners (org.apache.nifi.util.TestRunners)2 AbstractWebSocketSession (org.apache.nifi.websocket.AbstractWebSocketSession)2