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