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

Example 6 with MockStateManager

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

the class TestEnforceOrder method testWaitOvertakeSkip.

@Test
public void testWaitOvertakeSkip() 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.setProperty(EnforceOrder.MAX_ORDER, "10");
    runner.assertValid();
    Ordered.enqueue(runner, "b", 1);
    Ordered.enqueue(runner, "a", 2);
    Ordered.enqueue(runner, "a", 1);
    // waits for a.3 and a.4
    Ordered.enqueue(runner, "a", 5);
    // waits for b.2
    Ordered.enqueue(runner, "b", 3);
    // waits for c.1 to 8
    Ordered.enqueue(runner, "c", 9);
    // waits for d.1 to 9
    Ordered.enqueue(runner, "d", 10);
    runner.run();
    List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(EnforceOrder.REL_SUCCESS);
    assertEquals(3, succeeded.size());
    final FirstInFirstOutPrioritizer fifo = new FirstInFirstOutPrioritizer();
    succeeded.sort(fifo);
    succeeded.get(0).assertContentEquals("a.1");
    succeeded.get(1).assertContentEquals("a.2");
    succeeded.get(2).assertContentEquals("b.1");
    List<MockFlowFile> waiting = runner.getFlowFilesForRelationship(EnforceOrder.REL_WAIT);
    assertEquals(4, waiting.size());
    waiting.get(0).assertContentEquals("a.5");
    waiting.get(1).assertContentEquals("b.3");
    waiting.get(2).assertContentEquals("c.9");
    waiting.get(3).assertContentEquals("d.10");
    waiting.get(0).assertAttributeExists("EnforceOrder.startedAt");
    waiting.get(1).assertAttributeExists("EnforceOrder.startedAt");
    waiting.get(2).assertAttributeExists("EnforceOrder.startedAt");
    waiting.get(3).assertAttributeExists("EnforceOrder.startedAt");
    final MockStateManager stateManager = runner.getStateManager();
    stateManager.assertStateEquals("a.target", "3", Scope.LOCAL);
    stateManager.assertStateEquals("b.target", "2", Scope.LOCAL);
    stateManager.assertStateEquals("c.target", "1", Scope.LOCAL);
    stateManager.assertStateEquals("d.target", "1", Scope.LOCAL);
    stateManager.assertStateSet("a.updatedAt", Scope.LOCAL);
    stateManager.assertStateSet("b.updatedAt", Scope.LOCAL);
    stateManager.assertStateSet("c.updatedAt", Scope.LOCAL);
    stateManager.assertStateSet("d.updatedAt", Scope.LOCAL);
    // Run it again with waiting files.
    runner.clearTransferState();
    waiting.forEach(runner::enqueue);
    runner.run();
    runner.assertAllFlowFilesTransferred(EnforceOrder.REL_WAIT, 4);
    waiting = runner.getFlowFilesForRelationship(EnforceOrder.REL_WAIT);
    // Run it again with shorter wait timeout to make overtaking happen.
    runner.clearTransferState();
    runner.setProperty(EnforceOrder.WAIT_TIMEOUT, "10 ms");
    Thread.sleep(20);
    waiting.forEach(runner::enqueue);
    // arrived in time
    Ordered.enqueue(runner, "b", 2);
    // a.4 and a.5 have not arrived yet
    Ordered.enqueue(runner, "a", 6);
    runner.run();
    succeeded = runner.getFlowFilesForRelationship(EnforceOrder.REL_SUCCESS);
    succeeded.sort(fifo);
    assertEquals(3, succeeded.size());
    // This is ok because a.5 was there.
    succeeded.get(0).assertContentEquals("a.6");
    succeeded.get(1).assertContentEquals("b.2");
    succeeded.get(2).assertContentEquals("b.3");
    List<MockFlowFile> overtook = runner.getFlowFilesForRelationship(EnforceOrder.REL_OVERTOOK);
    assertEquals(3, overtook.size());
    // overtook a.3.
    overtook.get(0).assertContentEquals("a.5");
    overtook.get(0).assertAttributeEquals(EnforceOrder.ATTR_EXPECTED_ORDER, "3");
    // overtook c.1 - 8.
    overtook.get(1).assertContentEquals("c.9");
    overtook.get(1).assertAttributeEquals(EnforceOrder.ATTR_EXPECTED_ORDER, "1");
    // overtook d.1 - 9.
    overtook.get(2).assertContentEquals("d.10");
    overtook.get(2).assertAttributeEquals(EnforceOrder.ATTR_EXPECTED_ORDER, "1");
    stateManager.assertStateEquals("a.target", "7", Scope.LOCAL);
    stateManager.assertStateEquals("b.target", "4", Scope.LOCAL);
    // it was c.9, so +1
    stateManager.assertStateEquals("c.target", "10", Scope.LOCAL);
    // it was d.10 (max) so don't +1
    stateManager.assertStateEquals("d.target", "10", Scope.LOCAL);
    // Simulate a.3 and a.4 arrive but too late..
    runner.clearTransferState();
    Ordered.enqueue(runner, "a", 3);
    Ordered.enqueue(runner, "a", 4);
    runner.run();
    runner.assertAllFlowFilesTransferred(EnforceOrder.REL_SKIPPED, 2);
    final List<MockFlowFile> skipped = runner.getFlowFilesForRelationship(EnforceOrder.REL_SKIPPED);
    skipped.get(0).assertContentEquals("a.3");
    skipped.get(0).assertAttributeExists(EnforceOrder.ATTR_DETAIL);
    skipped.get(1).assertContentEquals("a.4");
    skipped.get(1).assertAttributeExists(EnforceOrder.ATTR_DETAIL);
}

Example 7 with MockStateManager

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

the class TestUpdateAttribute method testStateFailuresWithRulesUsingClone.

@Test
public void testStateFailuresWithRulesUsingClone() throws Exception {
    final Criteria criteria = getCriteria();
    criteria.setFlowFilePolicy(FlowFilePolicy.USE_CLONE);
    addRule(criteria, "rule", Collections.singletonList(// conditions
    "${getStateValue('maxValue'):lt(${value})}"), getMap(// actions
    "maxValue", "${value}"));
    addRule(criteria, "rule2", Collections.singletonList(// conditions
    "${getStateValue('maxValue2'):lt(${value})}"), getMap(// actions
    "maxValue2", "${value}"));
    TestRunner runner = TestRunners.newTestRunner(new UpdateAttribute());
    final UpdateAttribute processor = (UpdateAttribute) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    MockStateManager mockStateManager = runner.getStateManager();
    runner.setProperty(UpdateAttribute.STORE_STATE, STORE_STATE_LOCALLY);
    runner.setProperty(UpdateAttribute.STATEFUL_VARIABLES_INIT_VALUE, "0");
    runner.setAnnotationData(serialize(criteria));
    processor.onScheduled(runner.getProcessContext());
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("value", "1");
    runner.enqueue(new byte[0], attributes);
    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);
    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());
    runner.assertQueueNotEmpty();
    mockStateManager.setFailOnStateGet(Scope.LOCAL, false);
    mockStateManager.setFailOnStateSet(Scope.LOCAL, true);
    processor.onTrigger(runner.getProcessContext(), processSessionFactory.createSession());
    runner.assertQueueEmpty();
    runner.assertAllFlowFilesTransferred(UpdateAttribute.REL_FAILED_SET_STATE, 1);
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeEquals("maxValue", "1");
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeNotExists("maxValue2");
}

Example 8 with MockStateManager

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

the class TestTailFile method testMigrateFrom100To110FileNotFound.

@Test
public void testMigrateFrom100To110FileNotFound() throws IOException {
    assumeFalse(isWindowsEnvironment());
    runner.setProperty(TailFile.FILENAME, "target/not-existing-log.txt");
    final MockStateManager stateManager = runner.getStateManager();
    // Before NiFi 1.1.0, TailFile only handles single file
    // and state key doesn't have index in it.
    final Map<String, String> state = new HashMap<>();
    state.put("filename", "target/not-existing-log.txt");
    // Simulate that it has been tailed up to the 2nd line.
    state.put("checksum", "2279929157");
    state.put("position", "14");
    state.put("timestamp", "1480639134000");
    stateManager.setState(state, Scope.LOCAL);
    runner.run();
    runner.assertTransferCount(TailFile.REL_SUCCESS, 0);
}

Example 9 with MockStateManager

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

the class TestTailFile method testMigrateFrom100To110.

@Test
public void testMigrateFrom100To110() throws IOException {
    assumeFalse(isWindowsEnvironment());
    runner.setProperty(TailFile.FILENAME, "target/existing-log.txt");
    final MockStateManager stateManager = runner.getStateManager();
    // Before NiFi 1.1.0, TailFile only handles single file
    // and state key doesn't have index in it.
    final Map<String, String> state = new HashMap<>();
    state.put("filename", "target/existing-log.txt");
    // Simulate that it has been tailed up to the 2nd line.
    state.put("checksum", "2279929157");
    state.put("position", "14");
    state.put("timestamp", "1480639134000");
    stateManager.setState(state, Scope.LOCAL);
    runner.run();
    runner.assertAllFlowFilesTransferred(TailFile.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(TailFile.REL_SUCCESS).iterator().next();
    final ByteArrayOutputStream bos = new ByteArrayOutputStream();
    try (final BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(bos))) {
        writer.write("Line 3");
        writer.newLine();
    }
    flowFile.assertContentEquals(bos.toByteArray());
    // The old states should be replaced with new ones.
    final StateMap updatedState = stateManager.getState(Scope.LOCAL);
    assertNull(updatedState.get("filename"));
    assertNull(updatedState.get("checksum"));
    assertNull(updatedState.get("position"));
    assertNull(updatedState.get("timestamp"));
    assertEquals("target/existing-log.txt", updatedState.get("file.0.filename"));
    assertEquals("3380848603", updatedState.get("file.0.checksum"));
    assertEquals("21", updatedState.get("file.0.position"));
    assertNotNull(updatedState.get("file.0.timestamp"));
    // When it runs again, the state is already migrated, so it shouldn't emit any flow files.
    runner.clearTransferState();
    runner.run();
    runner.assertAllFlowFilesTransferred(TailFile.REL_SUCCESS, 0);
}

Example 10 with MockStateManager

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

the class TestAbstractListProcessor method testStateMigratedFromCacheService.

@Test
public void testStateMigratedFromCacheService() throws InitializationException {
    final DistributedCache cache = new DistributedCache();
    runner.addControllerService("cache", cache);
    runner.enableControllerService(cache);
    runner.setProperty(AbstractListProcessor.DISTRIBUTED_CACHE_SERVICE, "cache");
    final String serviceState = "{\"latestTimestamp\":1492,\"matchingIdentifiers\":[\"id\"]}";
    final String cacheKey = runner.getProcessor().getIdentifier() + ".lastListingTime./path";
    cache.stored.put(cacheKey, serviceState);
    runner.run();
    final MockStateManager stateManager = runner.getStateManager();
    final Map<String, String> expectedState = new HashMap<>();
    // Ensure only timestamp is migrated
    expectedState.put(AbstractListProcessor.LATEST_LISTED_ENTRY_TIMESTAMP_KEY, "1492");
    expectedState.put(AbstractListProcessor.LAST_PROCESSED_LATEST_ENTRY_TIMESTAMP_KEY, "1492");
    stateManager.assertStateEquals(expectedState, Scope.CLUSTER);
}