Examples with MockStateManager org.apache.nifi.state.MockStateManager used on opensource projects

Example 11 with MockStateManager

use of org.apache.nifi.state.MockStateManager in project nifi by apache.

the class TestAbstractListProcessor method testNoStateToMigrate.

@Test
public void testNoStateToMigrate() throws Exception {
    runner.run();
    final MockStateManager stateManager = runner.getStateManager();
    final Map<String, String> expectedState = new HashMap<>();
    stateManager.assertStateEquals(expectedState, Scope.CLUSTER);
}

Example 12 with MockStateManager

use of org.apache.nifi.state.MockStateManager in project nifi by apache.

the class TestSiteToSiteBulletinReportingTask method testSerializedForm.

@Test
public void testSerializedForm() throws IOException, InitializationException {
    // creating the list of bulletins
    final List<Bulletin> bulletins = new ArrayList<Bulletin>();
    bulletins.add(BulletinFactory.createBulletin("group-id", "group-name", "source-id", "source-name", "category", "severity", "message"));
    // mock the access to the list of bulletins
    final ReportingContext context = Mockito.mock(ReportingContext.class);
    final BulletinRepository repository = Mockito.mock(BulletinRepository.class);
    Mockito.when(context.getBulletinRepository()).thenReturn(repository);
    Mockito.when(repository.findBulletins(Mockito.any(BulletinQuery.class))).thenReturn(bulletins);
    // creating reporting task
    final MockSiteToSiteBulletinReportingTask task = new MockSiteToSiteBulletinReportingTask();
    Mockito.when(context.getStateManager()).thenReturn(new MockStateManager(task));
    // settings properties and mocking access to properties
    final Map<PropertyDescriptor, String> properties = new HashMap<>();
    for (final PropertyDescriptor descriptor : task.getSupportedPropertyDescriptors()) {
        properties.put(descriptor, descriptor.getDefaultValue());
    }
    properties.put(SiteToSiteBulletinReportingTask.BATCH_SIZE, "1000");
    properties.put(SiteToSiteBulletinReportingTask.PLATFORM, "nifi");
    Mockito.doAnswer(new Answer<PropertyValue>() {

        @Override
        public PropertyValue answer(final InvocationOnMock invocation) throws Throwable {
            final PropertyDescriptor descriptor = invocation.getArgumentAt(0, PropertyDescriptor.class);
            return new MockPropertyValue(properties.get(descriptor));
        }
    }).when(context).getProperty(Mockito.any(PropertyDescriptor.class));
    // setup the mock initialization context
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    final ReportingInitializationContext initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    Mockito.when(initContext.getLogger()).thenReturn(logger);
    task.initialize(initContext);
    task.onTrigger(context);
    // test checking
    assertEquals(1, task.dataSent.size());
    final String msg = new String(task.dataSent.get(0), StandardCharsets.UTF_8);
    JsonReader jsonReader = Json.createReader(new ByteArrayInputStream(msg.getBytes()));
    JsonObject bulletinJson = jsonReader.readArray().getJsonObject(0);
    assertEquals("message", bulletinJson.getString("bulletinMessage"));
    assertEquals("group-name", bulletinJson.getString("bulletinGroupName"));
}

Example 13 with MockStateManager

use of org.apache.nifi.state.MockStateManager in project nifi by apache.

the class TestSiteToSiteProvenanceReportingTask method testWhenProvenanceMaxIdEqualToLastEventIdInStateManager.

@Test
public void testWhenProvenanceMaxIdEqualToLastEventIdInStateManager() throws IOException, InitializationException {
    final long maxEventId = 2500;
    // create the mock reporting task and mock state manager
    final Map<PropertyDescriptor, String> properties = new HashMap<>();
    for (final PropertyDescriptor descriptor : new MockSiteToSiteProvenanceReportingTask().getSupportedPropertyDescriptors()) {
        properties.put(descriptor, descriptor.getDefaultValue());
    }
    final MockSiteToSiteProvenanceReportingTask task = setup(null, properties);
    final MockStateManager stateManager = new MockStateManager(task);
    // create the state map and set the last id to the same value as maxEventId
    final Map<String, String> state = new HashMap<>();
    state.put(SiteToSiteProvenanceReportingTask.LAST_EVENT_ID_KEY, String.valueOf(maxEventId));
    stateManager.setState(state, Scope.LOCAL);
    // setup the mock provenance repository to return maxEventId
    final ProvenanceEventRepository provenanceRepository = Mockito.mock(ProvenanceEventRepository.class);
    Mockito.doAnswer(new Answer<Long>() {

        @Override
        public Long answer(final InvocationOnMock invocation) throws Throwable {
            return maxEventId;
        }
    }).when(provenanceRepository).getMaxEventId();
    // setup the mock EventAccess to return the mock provenance repository
    final EventAccess eventAccess = Mockito.mock(EventAccess.class);
    when(eventAccess.getProvenanceRepository()).thenReturn(provenanceRepository);
    task.initialize(initContext);
    // execute the reporting task and should not produce any data b/c max id same as previous id
    task.onScheduled(confContext);
    task.onTrigger(context);
    assertEquals(0, task.dataSent.size());
}

Example 14 with MockStateManager

use of org.apache.nifi.state.MockStateManager in project nifi by apache.

the class TestSiteToSiteProvenanceReportingTask method setup.

private MockSiteToSiteProvenanceReportingTask setup(ProvenanceEventRecord event, Map<PropertyDescriptor, String> properties, long maxEventId) throws IOException {
    final MockSiteToSiteProvenanceReportingTask task = new MockSiteToSiteProvenanceReportingTask();
    when(context.getStateManager()).thenReturn(new MockStateManager(task));
    Mockito.doAnswer(new Answer<PropertyValue>() {

        @Override
        public PropertyValue answer(final InvocationOnMock invocation) throws Throwable {
            final PropertyDescriptor descriptor = invocation.getArgumentAt(0, PropertyDescriptor.class);
            return new MockPropertyValue(properties.get(descriptor));
        }
    }).when(context).getProperty(Mockito.any(PropertyDescriptor.class));
    Mockito.doAnswer(new Answer<PropertyValue>() {

        @Override
        public PropertyValue answer(final InvocationOnMock invocation) throws Throwable {
            final PropertyDescriptor descriptor = invocation.getArgumentAt(0, PropertyDescriptor.class);
            return new MockPropertyValue(properties.get(descriptor));
        }
    }).when(confContext).getProperty(Mockito.any(PropertyDescriptor.class));
    final AtomicInteger totalEvents = new AtomicInteger(0);
    final EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.doAnswer(new Answer<List<ProvenanceEventRecord>>() {

        @Override
        public List<ProvenanceEventRecord> answer(final InvocationOnMock invocation) throws Throwable {
            final long startId = invocation.getArgumentAt(0, long.class);
            final int maxRecords = invocation.getArgumentAt(1, int.class);
            final List<ProvenanceEventRecord> eventsToReturn = new ArrayList<>();
            for (int i = (int) Math.max(0, startId); i < (int) (startId + maxRecords) && totalEvents.get() < maxEventId; i++) {
                if (event != null) {
                    eventsToReturn.add(event);
                }
                totalEvents.getAndIncrement();
            }
            return eventsToReturn;
        }
    }).when(eventAccess).getProvenanceEvents(Mockito.anyLong(), Mockito.anyInt());
    ProcessGroupStatus pgRoot = new ProcessGroupStatus();
    pgRoot.setId("root");
    when(eventAccess.getControllerStatus()).thenReturn(pgRoot);
    // Add child Process Groups.
    // Root -> (A, B -> (B2 -> (B3)))
    final ProcessGroupStatus pgA = new ProcessGroupStatus();
    pgA.setId("pgA");
    final ProcessGroupStatus pgB = new ProcessGroupStatus();
    pgB.setId("pgB");
    final ProcessGroupStatus pgB2 = new ProcessGroupStatus();
    pgB2.setId("pgB2");
    final ProcessGroupStatus pgB3 = new ProcessGroupStatus();
    pgB3.setId("pgB3");
    final Collection<ProcessGroupStatus> childPGs = pgRoot.getProcessGroupStatus();
    childPGs.add(pgA);
    childPGs.add(pgB);
    pgB.getProcessGroupStatus().add(pgB2);
    pgB2.getProcessGroupStatus().add(pgB3);
    // Add Processors.
    final ProcessorStatus prcRoot = new ProcessorStatus();
    prcRoot.setId("1234");
    pgRoot.getProcessorStatus().add(prcRoot);
    final ProcessorStatus prcA = new ProcessorStatus();
    prcA.setId("A001");
    prcA.setName("Processor in PGA");
    pgA.getProcessorStatus().add(prcA);
    final ProcessorStatus prcB = new ProcessorStatus();
    prcB.setId("B001");
    prcB.setName("Processor in PGB");
    pgB.getProcessorStatus().add(prcB);
    final ProcessorStatus prcB2 = new ProcessorStatus();
    prcB2.setId("B201");
    prcB2.setName("Processor in PGB2");
    pgB2.getProcessorStatus().add(prcB2);
    final ProcessorStatus prcB3 = new ProcessorStatus();
    prcB3.setId("B301");
    prcB3.setName("Processor in PGB3");
    pgB3.getProcessorStatus().add(prcB3);
    // Add connection status to test Remote Input/Output Ports
    final ConnectionStatus b2RemoteInputPort = new ConnectionStatus();
    b2RemoteInputPort.setGroupId("pgB2");
    b2RemoteInputPort.setSourceId("B201");
    b2RemoteInputPort.setDestinationId("riB2");
    b2RemoteInputPort.setDestinationName("Remote Input Port name");
    pgB2.getConnectionStatus().add(b2RemoteInputPort);
    final ConnectionStatus b3RemoteOutputPort = new ConnectionStatus();
    b3RemoteOutputPort.setGroupId("pgB3");
    b3RemoteOutputPort.setSourceId("roB3");
    b3RemoteOutputPort.setSourceName("Remote Output Port name");
    b3RemoteOutputPort.setDestinationId("B301");
    pgB3.getConnectionStatus().add(b3RemoteOutputPort);
    final ProvenanceEventRepository provenanceRepository = Mockito.mock(ProvenanceEventRepository.class);
    Mockito.doAnswer(new Answer<Long>() {

        @Override
        public Long answer(final InvocationOnMock invocation) throws Throwable {
            return maxEventId;
        }
    }).when(provenanceRepository).getMaxEventId();
    when(context.getEventAccess()).thenReturn(eventAccess);
    when(eventAccess.getProvenanceRepository()).thenReturn(provenanceRepository);
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    when(initContext.getLogger()).thenReturn(logger);
    return task;
}

Example 15 with MockStateManager

use of org.apache.nifi.state.MockStateManager in project nifi by apache.

the class TestSiteToSiteStatusReportingTask method initTask.

public MockSiteToSiteStatusReportingTask initTask(Map<PropertyDescriptor, String> customProperties, ProcessGroupStatus pgStatus) throws InitializationException {
    final MockSiteToSiteStatusReportingTask task = new MockSiteToSiteStatusReportingTask();
    Map<PropertyDescriptor, String> properties = new HashMap<>();
    for (final PropertyDescriptor descriptor : task.getSupportedPropertyDescriptors()) {
        properties.put(descriptor, descriptor.getDefaultValue());
    }
    properties.putAll(customProperties);
    context = Mockito.mock(ReportingContext.class);
    Mockito.when(context.getStateManager()).thenReturn(new MockStateManager(task));
    Mockito.doAnswer(new Answer<PropertyValue>() {

        @Override
        public PropertyValue answer(final InvocationOnMock invocation) throws Throwable {
            final PropertyDescriptor descriptor = invocation.getArgumentAt(0, PropertyDescriptor.class);
            return new MockPropertyValue(properties.get(descriptor));
        }
    }).when(context).getProperty(Mockito.any(PropertyDescriptor.class));
    final EventAccess eventAccess = Mockito.mock(EventAccess.class);
    Mockito.when(context.getEventAccess()).thenReturn(eventAccess);
    Mockito.when(eventAccess.getControllerStatus()).thenReturn(pgStatus);
    final ComponentLog logger = Mockito.mock(ComponentLog.class);
    final ReportingInitializationContext initContext = Mockito.mock(ReportingInitializationContext.class);
    Mockito.when(initContext.getIdentifier()).thenReturn(UUID.randomUUID().toString());
    Mockito.when(initContext.getLogger()).thenReturn(logger);
    task.initialize(initContext);
    return task;
}