Search in sources :

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

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));
}
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 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())));
}
Also used : Mockito.spy(org.mockito.Mockito.spy) HashSet(java.util.HashSet) Relationship(org.apache.nifi.processor.Relationship) TestRunner(org.apache.nifi.util.TestRunner) ProvenanceEventRecord(org.apache.nifi.provenance.ProvenanceEventRecord) Map(java.util.Map) Mockito.doAnswer(org.mockito.Mockito.doAnswer) WebSocketSession(org.apache.nifi.websocket.WebSocketSession) MockProcessSession(org.apache.nifi.util.MockProcessSession) ProvenanceEventType(org.apache.nifi.provenance.ProvenanceEventType) WebSocketMessage(org.apache.nifi.websocket.WebSocketMessage) 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) AtomicLong(java.util.concurrent.atomic.AtomicLong) List(java.util.List) SharedSessionState(org.apache.nifi.util.SharedSessionState) AbstractWebSocketSession(org.apache.nifi.websocket.AbstractWebSocketSession) WebSocketClientService(org.apache.nifi.websocket.WebSocketClientService) TestRunners(org.apache.nifi.util.TestRunners) Assert.assertEquals(org.junit.Assert.assertEquals) MockFlowFile(org.apache.nifi.util.MockFlowFile) Mockito.mock(org.mockito.Mockito.mock) WebSocketClientService(org.apache.nifi.websocket.WebSocketClientService) TestRunner(org.apache.nifi.util.TestRunner) InetSocketAddress(java.net.InetSocketAddress) 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) Relationship(org.apache.nifi.processor.Relationship) ProvenanceEventRecord(org.apache.nifi.provenance.ProvenanceEventRecord) ProcessSessionFactory(org.apache.nifi.processor.ProcessSessionFactory) List(java.util.List) MockProcessSession(org.apache.nifi.util.MockProcessSession) HashSet(java.util.HashSet) Test(org.junit.Test)

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

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

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