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

Example 16 with MockStateManager

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

the class MetricsReportingTaskTest method setUp.

/**
 * Set up the test environment and mock behaviour. This includes registering {@link #reporterServiceStub} in the
 * different contexts, overriding {@link MetricsReportingTask#currentStatusReference} and instantiating the test
 * subject.
 */
@Before
public void setUp() throws Exception {
    Map<String, ControllerService> services = new HashMap<>();
    services.put(REPORTER_SERVICE_IDENTIFIER, reporterServiceStub);
    testedReportingTask = new MetricsReportingTask();
    reportingContextStub = new MockReportingContext(services, new MockStateManager(testedReportingTask), new MockVariableRegistry());
    rootGroupStatus = new ProcessGroupStatus();
    innerGroupStatus = new ProcessGroupStatus();
    when(reporterServiceStub.createReporter(any())).thenReturn(reporterMock);
    when(reporterServiceStub.getIdentifier()).thenReturn(REPORTER_SERVICE_IDENTIFIER);
    reportingContextStub.setProperty(MetricsReportingTask.REPORTER_SERVICE.getName(), REPORTER_SERVICE_IDENTIFIER);
    reportingContextStub.addControllerService(reporterServiceStub, REPORTER_SERVICE_IDENTIFIER);
    configurationContextStub = new MockConfigurationContext(reportingContextStub.getProperties(), reportingContextStub.getControllerServiceLookup());
    reportingInitContextStub = new MockReportingInitializationContext(TEST_INIT_CONTEXT_ID, TEST_INIT_CONTEXT_NAME, new MockComponentLog(TEST_TASK_ID, testedReportingTask));
}

Example 17 with MockStateManager

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

the class TestDetectDuplicate method createClient.

private DistributedMapCacheClientImpl createClient() throws InitializationException {
    final DistributedMapCacheClientImpl client = new DistributedMapCacheClientImpl();
    final ComponentLog logger = new MockComponentLog("client", client);
    final MockControllerServiceInitializationContext clientInitContext = new MockControllerServiceInitializationContext(client, "client", logger, new MockStateManager(client));
    client.initialize(clientInitContext);
    return client;
}

Example 18 with MockStateManager

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

the class TestEnforceOrder method testCleanInactiveGroups.

@Test
public void testCleanInactiveGroups() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(EnforceOrder.class);
    runner.setProperty(EnforceOrder.GROUP_IDENTIFIER, "${group}");
    runner.setProperty(EnforceOrder.ORDER_ATTRIBUTE, "index");
    runner.setProperty(EnforceOrder.INITIAL_ORDER, "1");
    runner.assertValid();
    Ordered.enqueue(runner, "b", 1);
    Ordered.enqueue(runner, "a", 2);
    Ordered.enqueue(runner, "c", 1);
    Ordered.enqueue(runner, "a", 1);
    runner.run();
    runner.assertAllFlowFilesTransferred(EnforceOrder.REL_SUCCESS, 4);
    // Run it again with shorter inactive timeout
    runner.clearTransferState();
    runner.setProperty(EnforceOrder.WAIT_TIMEOUT, "5 ms");
    runner.setProperty(EnforceOrder.INACTIVE_TIMEOUT, "10 ms");
    Thread.sleep(15);
    // No group b.
    Ordered.enqueue(runner, "a", 3);
    Ordered.enqueue(runner, "c", 2);
    runner.run();
    // Group b was determined as inactive, thus its states should be removed.
    final MockStateManager stateManager = runner.getStateManager();
    stateManager.assertStateEquals("a.target", "4", Scope.LOCAL);
    stateManager.assertStateNotSet("b.target", Scope.LOCAL);
    stateManager.assertStateEquals("c.target", "3", Scope.LOCAL);
    stateManager.assertStateSet("a.updatedAt", Scope.LOCAL);
    stateManager.assertStateNotSet("b.updatedAt", Scope.LOCAL);
    stateManager.assertStateSet("c.updatedAt", Scope.LOCAL);
    // If b comes again, it'll be treated as brand new group.
    runner.clearTransferState();
    Ordered.enqueue(runner, "b", 2);
    runner.run();
    stateManager.assertStateEquals("b.target", "1", Scope.LOCAL);
    stateManager.assertStateSet("b.updatedAt", Scope.LOCAL);
    // b.2 should be routed to wait, since there's no b.1. It will eventually overtake.
    runner.assertAllFlowFilesTransferred(EnforceOrder.REL_WAIT, 1);
    final List<MockFlowFile> waiting = runner.getFlowFilesForRelationship(EnforceOrder.REL_WAIT);
    waiting.get(0).assertContentEquals("b.2");
}

Example 19 with MockStateManager

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

the class TestEnforceOrder method testInitialOrderValue.

@Test
public void testInitialOrderValue() {
    final TestRunner runner = TestRunners.newTestRunner(EnforceOrder.class);
    runner.setProperty(EnforceOrder.GROUP_IDENTIFIER, "${group}");
    runner.setProperty(EnforceOrder.ORDER_ATTRIBUTE, "index");
    runner.setProperty(EnforceOrder.INITIAL_ORDER, "${index.start}");
    runner.setProperty(EnforceOrder.MAX_ORDER, "${index.max}");
    runner.assertValid();
    runner.enqueue("b.0", Ordered.i("b", 0).put("index.start", "0").put("index.max", "99").map());
    runner.enqueue("a.100", Ordered.i("a", 100).put("index.start", "100").put("index.max", "103").map());
    runner.enqueue("a.101", Ordered.i("a", 101).put("index.start", "100").put("index.max", "103").map());
    runner.enqueue("illegal initial order", Ordered.i("c", 1).put("index.start", "non-integer").map());
    runner.enqueue("without initial order", Ordered.i("d", 1).map());
    // Even if this flow file doesn't have initial order attribute, this will be routed to success.
    // Because target order for group b is already computed from b.0.
    Ordered.enqueue(runner, "b", 1);
    runner.run();
    List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(EnforceOrder.REL_SUCCESS);
    assertEquals(4, succeeded.size());
    succeeded.sort(new FirstInFirstOutPrioritizer());
    succeeded.get(0).assertContentEquals("a.100");
    succeeded.get(1).assertContentEquals("a.101");
    succeeded.get(2).assertContentEquals("b.0");
    succeeded.get(3).assertContentEquals("b.1");
    final List<MockFlowFile> failed = runner.getFlowFilesForRelationship(EnforceOrder.REL_FAILURE);
    assertEquals(2, failed.size());
    failed.get(0).assertAttributeExists(EnforceOrder.ATTR_DETAIL);
    failed.get(0).assertContentEquals("illegal initial order");
    failed.get(1).assertAttributeExists(EnforceOrder.ATTR_DETAIL);
    failed.get(1).assertContentEquals("without initial order");
    final MockStateManager stateManager = runner.getStateManager();
    stateManager.assertStateEquals("a.target", "102", Scope.LOCAL);
    stateManager.assertStateEquals("a.max", "103", Scope.LOCAL);
    stateManager.assertStateEquals("b.target", "2", Scope.LOCAL);
    stateManager.assertStateEquals("b.max", "99", Scope.LOCAL);
    runner.clearTransferState();
}

Example 20 with MockStateManager

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

the class ITReportLineageToAtlas method test.

private void test(TestConfiguration tc) throws InitializationException, IOException {
    final ReportLineageToAtlas reportingTask = new ReportLineageToAtlas();
    final MockComponentLog logger = new MockComponentLog("reporting-task-id", reportingTask);
    final ReportingInitializationContext initializationContext = mock(ReportingInitializationContext.class);
    when(initializationContext.getLogger()).thenReturn(logger);
    final ConfigurationContext configurationContext = new MockConfigurationContext(tc.properties, null);
    final ValidationContext validationContext = mock(ValidationContext.class);
    when(validationContext.getProperty(any())).then(invocation -> new MockPropertyValue(tc.properties.get(invocation.getArguments()[0])));
    final ReportingContext reportingContext = mock(ReportingContext.class);
    final MockStateManager stateManager = new MockStateManager(reportingTask);
    final EventAccess eventAccess = mock(EventAccess.class);
    when(reportingContext.getProperties()).thenReturn(tc.properties);
    when(reportingContext.getProperty(any())).then(invocation -> new MockPropertyValue(tc.properties.get(invocation.getArguments()[0])));
    when(reportingContext.getStateManager()).thenReturn(stateManager);
    when(reportingContext.getEventAccess()).thenReturn(eventAccess);
    when(eventAccess.getGroupStatus(eq("root"))).thenReturn(tc.rootPgStatus);
    final ProvenanceRepository provenanceRepository = mock(ProvenanceRepository.class);
    when(eventAccess.getControllerStatus()).thenReturn(tc.rootPgStatus);
    when(eventAccess.getProvenanceRepository()).thenReturn(provenanceRepository);
    when(eventAccess.getProvenanceEvents(eq(-1L), anyInt())).thenReturn(tc.provenanceRecords);
    when(provenanceRepository.getMaxEventId()).thenReturn((long) tc.provenanceRecords.size() - 1);
    when(provenanceRepository.getEvent(anyLong())).then(invocation -> tc.provenanceRecords.get(((Long) invocation.getArguments()[0]).intValue()));
    // To mock this async method invocations, keep the requested event ids in a stack.
    final ComputeLineageSubmission lineageComputationSubmission = mock(ComputeLineageSubmission.class);
    when(provenanceRepository.submitLineageComputation(anyLong(), any())).thenAnswer(invocation -> {
        requestedLineageComputationIds.push((Long) invocation.getArguments()[0]);
        return lineageComputationSubmission;
    });
    when(lineageComputationSubmission.getResult()).then(invocation -> tc.lineageResults.get(requestedLineageComputationIds.pop()));
    final ComputeLineageSubmission expandParentsSubmission = mock(ComputeLineageSubmission.class);
    when(provenanceRepository.submitExpandParents(anyLong(), any())).thenAnswer(invocation -> {
        requestedExpandParentsIds.push(((Long) invocation.getArguments()[0]));
        return expandParentsSubmission;
    });
    when(expandParentsSubmission.getResult()).then(invocation -> tc.parentLineageResults.get(requestedExpandParentsIds.pop()));
    tc.properties.put(ATLAS_NIFI_URL, "http://localhost:8080/nifi");
    tc.properties.put(ATLAS_URLS, TARGET_ATLAS_URL);
    tc.properties.put(ATLAS_USER, "admin");
    tc.properties.put(ATLAS_PASSWORD, "admin");
    tc.properties.put(new PropertyDescriptor.Builder().name("hostnamePattern.example").dynamic(true).build(), ".*");
    reportingTask.initialize(initializationContext);
    reportingTask.validate(validationContext);
    reportingTask.setup(configurationContext);
    reportingTask.onTrigger(reportingContext);
    reportingTask.onUnscheduled();
    reportingTask.onStopped();
}