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

Example 1 with MockStateManager

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

the class TestListS3 method testListObjectsNothingNew.

@Test
public void testListObjectsNothingNew() throws IOException {
    runner.setProperty(ListS3.REGION, "eu-west-1");
    runner.setProperty(ListS3.BUCKET, "test-bucket");
    Calendar calendar = Calendar.getInstance();
    calendar.set(2017, 5, 2);
    Date objectLastModified = calendar.getTime();
    long stateCurrentTimestamp = objectLastModified.getTime();
    Map<String, String> state = new HashMap<>();
    state.put(ListS3.CURRENT_TIMESTAMP, String.valueOf(stateCurrentTimestamp));
    state.put(ListS3.CURRENT_KEY_PREFIX + "0", "test-key");
    MockStateManager mockStateManager = runner.getStateManager();
    mockStateManager.setState(state, Scope.CLUSTER);
    ObjectListing objectListing = new ObjectListing();
    S3ObjectSummary objectSummary1 = new S3ObjectSummary();
    objectSummary1.setBucketName("test-bucket");
    objectSummary1.setKey("test-key");
    objectSummary1.setLastModified(objectLastModified);
    objectListing.getObjectSummaries().add(objectSummary1);
    Mockito.when(mockS3Client.listObjects(Mockito.any(ListObjectsRequest.class))).thenReturn(objectListing);
    runner.run();
    ArgumentCaptor<ListObjectsRequest> captureRequest = ArgumentCaptor.forClass(ListObjectsRequest.class);
    Mockito.verify(mockS3Client, Mockito.times(1)).listObjects(captureRequest.capture());
    ListObjectsRequest request = captureRequest.getValue();
    assertEquals("test-bucket", request.getBucketName());
    Mockito.verify(mockS3Client, Mockito.never()).listVersions(Mockito.any());
    runner.assertAllFlowFilesTransferred(ListS3.REL_SUCCESS, 0);
}

Example 2 with MockStateManager

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

the class StandardProcessorTestRunner method addControllerService.

@Override
public void addControllerService(final String identifier, final ControllerService service, final Map<String, String> properties) throws InitializationException {
    final MockComponentLog logger = new MockComponentLog(identifier, service);
    controllerServiceLoggers.put(identifier, logger);
    final MockStateManager serviceStateManager = new MockStateManager(service);
    final MockControllerServiceInitializationContext initContext = new MockControllerServiceInitializationContext(requireNonNull(service), requireNonNull(identifier), logger, serviceStateManager);
    controllerServiceStateManagers.put(identifier, serviceStateManager);
    initContext.addControllerServices(context);
    service.initialize(initContext);
    final Map<PropertyDescriptor, String> resolvedProps = new HashMap<>();
    for (final Map.Entry<String, String> entry : properties.entrySet()) {
        resolvedProps.put(service.getPropertyDescriptor(entry.getKey()), entry.getValue());
    }
    try {
        ReflectionUtils.invokeMethodsWithAnnotation(OnAdded.class, service);
    } catch (final InvocationTargetException | IllegalAccessException | IllegalArgumentException e) {
        throw new InitializationException(e);
    }
    context.addControllerService(identifier, service, resolvedProps, null);
}

Example 3 with MockStateManager

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

the class StandardProcessorTestRunner method setProperty.

@Override
public ValidationResult setProperty(final ControllerService service, final PropertyDescriptor property, final String value) {
    final MockStateManager serviceStateManager = controllerServiceStateManagers.get(service.getIdentifier());
    if (serviceStateManager == null) {
        throw new IllegalStateException("Controller service " + service + " has not been added to this TestRunner via the #addControllerService method");
    }
    final ControllerServiceConfiguration configuration = getConfigToUpdate(service);
    final Map<PropertyDescriptor, String> curProps = configuration.getProperties();
    final Map<PropertyDescriptor, String> updatedProps = new HashMap<>(curProps);
    final ValidationContext validationContext = new MockValidationContext(context, serviceStateManager, variableRegistry).getControllerServiceValidationContext(service);
    final ValidationResult validationResult = property.validate(value, validationContext);
    final String oldValue = updatedProps.get(property);
    updatedProps.put(property, value);
    configuration.setProperties(updatedProps);
    if ((value == null && oldValue != null) || (value != null && !value.equals(oldValue))) {
        service.onPropertyModified(property, oldValue, value);
    }
    return validationResult;
}

Example 4 with MockStateManager

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

the class TestSiteToSiteBulletinReportingTask method testWhenProvenanceMaxIdEqualToLastEventIdInStateManager.

@Test
public void testWhenProvenanceMaxIdEqualToLastEventIdInStateManager() throws IOException, InitializationException {
    // creating the list of bulletins
    final List<Bulletin> bulletins = new ArrayList<Bulletin>();
    bulletins.add(BulletinFactory.createBulletin("category", "severity", "message"));
    bulletins.add(BulletinFactory.createBulletin("category", "severity", "message"));
    bulletins.add(BulletinFactory.createBulletin("category", "severity", "message"));
    bulletins.add(BulletinFactory.createBulletin("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);
    final long maxEventId = getMaxBulletinId(bulletins);
    ;
    // create the mock reporting task and mock state manager
    final MockSiteToSiteBulletinReportingTask task = new MockSiteToSiteBulletinReportingTask();
    final MockStateManager stateManager = 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");
    properties.put(SiteToSiteBulletinReportingTask.TRANSPORT_PROTOCOL, SiteToSiteTransportProtocol.HTTP.name());
    properties.put(SiteToSiteBulletinReportingTask.HTTP_PROXY_HOSTNAME, "localhost");
    properties.put(SiteToSiteBulletinReportingTask.HTTP_PROXY_PORT, "80");
    properties.put(SiteToSiteBulletinReportingTask.HTTP_PROXY_USERNAME, "username");
    properties.put(SiteToSiteBulletinReportingTask.HTTP_PROXY_PASSWORD, "password");
    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));
    // 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 reporting context to return the mock state manager
    Mockito.when(context.getStateManager()).thenReturn(stateManager);
    // 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);
    // execute the reporting task and should not produce any data b/c max id same as previous id
    task.onTrigger(context);
    assertEquals(0, task.dataSent.size());
}

Example 5 with MockStateManager

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

the class TestEnforceOrder method testMaxOrder.

@Test
public void testMaxOrder() {
    final TestRunner runner = TestRunners.newTestRunner(EnforceOrder.class);
    runner.setProperty(EnforceOrder.GROUP_IDENTIFIER, "${fragment.identifier}");
    runner.setProperty(EnforceOrder.ORDER_ATTRIBUTE, "index");
    runner.setProperty(EnforceOrder.INITIAL_ORDER, "1");
    runner.setProperty(EnforceOrder.MAX_ORDER, "${fragment.count}");
    runner.assertValid();
    runner.enqueue("b.1", Ordered.i(1).put("fragment.identifier", "b").put("fragment.count", "3").map());
    runner.enqueue("a.2", Ordered.i(2).put("fragment.identifier", "a").put("fragment.count", "2").map());
    runner.enqueue("without max order", Ordered.i(1).put("fragment.identifier", "c").map());
    runner.enqueue("illegal max order", Ordered.i(1).put("fragment.identifier", "d").put("fragment.count", "X").map());
    runner.enqueue("a.1", Ordered.i(1).put("fragment.identifier", "a").put("fragment.count", "2").map());
    // Exceed max
    runner.enqueue("a.3", Ordered.i(3).put("fragment.identifier", "a").put("fragment.count", "2").map());
    runner.run();
    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(EnforceOrder.REL_SUCCESS);
    succeeded.sort(new FirstInFirstOutPrioritizer());
    assertEquals(3, succeeded.size());
    succeeded.get(0).assertContentEquals("a.1");
    succeeded.get(1).assertContentEquals("a.2");
    succeeded.get(2).assertContentEquals("b.1");
    final List<MockFlowFile> failed = runner.getFlowFilesForRelationship(EnforceOrder.REL_FAILURE);
    assertEquals(3, failed.size());
    failed.get(0).assertContentEquals("without max order");
    failed.get(1).assertContentEquals("illegal max order");
    // exceeds max order
    failed.get(2).assertContentEquals("a.3");
    final MockStateManager stateManager = runner.getStateManager();
    stateManager.assertStateEquals("a.target", "2", Scope.LOCAL);
    stateManager.assertStateEquals("a.max", "2", Scope.LOCAL);
}