Java Code Examples for org.apache.nifi.util.TestRunner#assertQueueNotEmpty()

@Test
public void testStateFailures() throws Exception {
    final 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("count", "${getStateValue('count'):plus(1)}");
    runner.setProperty("sum", "${getStateValue('sum'):plus(${pencils})}");
    runner.setProperty(UpdateAttribute.STATEFUL_VARIABLES_INIT_VALUE, "0");

    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes2 = new HashMap<>();
    attributes2.put("pencils", "2");

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);

    runner.enqueue(new byte[0],attributes2);
    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("count", "1");
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_FAILED_SET_STATE).get(0).assertAttributeEquals("sum", "2");
}
 

/**
 * This test will verify that if we have a FlowFile larger than the Max Size for a Bin, it will go into its
 * own bin and immediately be processed as its own bin.
 */
@Test
public void testFlowFileLargerThanBin() {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MERGE_STRATEGY, MergeContent.MERGE_STRATEGY_BIN_PACK);
    runner.setProperty(MergeContent.MIN_ENTRIES, "2");
    runner.setProperty(MergeContent.MAX_ENTRIES, "2");
    runner.setProperty(MergeContent.MIN_SIZE, "1 KB");
    runner.setProperty(MergeContent.MAX_SIZE, "5 KB");

    runner.enqueue(new byte[1026]); // add flowfile that fits within the bin limits
    runner.enqueue(new byte[1024 * 6]); // add flowfile that is larger than the bin limit
    runner.run(2); // run twice so that we have a chance to create two bins (though we shouldn't create 2, because only 1 bin will be full)

    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 1);
    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    assertEquals(runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0).getAttribute(CoreAttributes.UUID.key()),
            runner.getFlowFilesForRelationship(MergeContent.REL_ORIGINAL).get(0).getAttribute(MergeContent.MERGE_UUID_ATTRIBUTE));

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0);
    assertEquals(1024 * 6, bundle.getSize());

    // Queue should not be empty because the first FlowFile will be transferred back to the input queue
    // when we run out @OnStopped logic, since it won't be transferred to any bin.
    runner.assertQueueNotEmpty();
}
 

@Test
public void testBatchLimit() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.FLOWFILE_RATE);
    runner.setProperty(ControlRate.MAX_RATE, "5555");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");

    final int TEST_FILE_COUNT = 1500;

    for (int i = 0; i < TEST_FILE_COUNT; i++) {
        runner.enqueue("test data " + i);
    }

    runner.run(1, false);

    // after 1 run should have MAX_FLOW_FILES_PER_BATCH files transferred and remainder of TEST_FILE_COUNT in queue
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, MAX_FLOW_FILES_PER_BATCH);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();
    assertEquals(TEST_FILE_COUNT - MAX_FLOW_FILES_PER_BATCH, runner.getQueueSize().getObjectCount());

    runner.run(1, false);

    // after 2 runs should have TEST_FILE_COUNT files transferred and 0 in queue
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, TEST_FILE_COUNT);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueEmpty();
}
 

@Test
public void testViaAttribute() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.ATTRIBUTE_RATE);
    runner.setProperty(ControlRate.RATE_CONTROL_ATTRIBUTE_NAME, "count");
    runner.setProperty(ControlRate.MAX_RATE, "20000");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");

    createFlowFile(runner, 1000);
    createFlowFile(runner, 3000);
    createFlowFile(runner, 5000);
    createFlowFile(runner, 20000);
    createFlowFile(runner, 1000);

    runner.run(5, false);

    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 4);
    runner.clearTransferState();

    // at this point, we have sent through 29,000 but our max is 20,000 per second.
    // After 1.45 seconds (29000 / 20000), we should be able to send another 20,000
    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();
    Thread.sleep(1200L);

    // at this point, more than TIME_PERIOD 1.0 seconds but less than 1.45 seconds have passed
    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();
    Thread.sleep(600L);

    // at this point, more than 1.45 seconds have passed, so we should be able to send another 20,000
    runner.run();
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 1);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueEmpty();
}
 

@Test
public void testDataSizeRate() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.DATA_RATE);
    runner.setProperty(ControlRate.MAX_RATE, "20 b");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");

    runner.enqueue("testdata 1");
    runner.enqueue("testdata 2");
    runner.enqueue("testdata 3");
    runner.enqueue("testdata 4");

    runner.run(4, false);

    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 2);
    runner.clearTransferState();

    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();

    // we have sent 20 bytes and after 1 second, we should be able to send 20 more
    Thread.sleep(1100L);
    runner.run(2, false);
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 2);
    runner.assertQueueEmpty();
}
 

@Test
public void testFileCountWithGrouping() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.FLOWFILE_RATE);
    runner.setProperty(ControlRate.MAX_RATE, "2");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");
    runner.setProperty(ControlRate.GROUPING_ATTRIBUTE_NAME, "group");

    createFlowFileWithGroup(runner, "one");
    createFlowFileWithGroup(runner, "two");
    createFlowFileWithGroup(runner, "one");
    createFlowFileWithGroup(runner, "two");
    createFlowFileWithGroup(runner, "one");
    createFlowFileWithGroup(runner, "two");

    runner.run(6, false);

    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 4);
    runner.clearTransferState();

    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();

    // we have sent 2 files per group and after 1 second, we should be able to send the remaining 1 file per group
    Thread.sleep(1100L);
    runner.run(2);
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 2);
    runner.assertQueueEmpty();
}
 

@Test
public void testFileCountRate() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.FLOWFILE_RATE);
    runner.setProperty(ControlRate.MAX_RATE, "3");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");

    runner.enqueue("test data 1");
    runner.enqueue("test data 2");
    runner.enqueue("test data 3");
    runner.enqueue("test data 4");

    runner.run(4, false);

    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 3);
    runner.clearTransferState();

    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();

    // we have sent 3 files and after 1 second, we should be able to send the 4th
    Thread.sleep(1100L);
    runner.run();
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 1);
    runner.assertQueueEmpty();
}
 

@Test(expected = AssertionError.class)
public void testFailedDeref() {
  TestRunner runner = getTestRunner(true);
  MockFlowFile flowFile = new MockFlowFile(1);
  Map<String, String> attributes = new HashMap<>();
  attributes.put("riemann.metric", "5");
  flowFile.putAttributes(attributes);
  runner.enqueue(flowFile);
  try {
    runner.run();
  } catch (ProcessException e) {
    runner.assertQueueNotEmpty();
    throw e;
  }
}
 

@Test(expected = AssertionError.class)
public void testFailedDeref() {
  TestRunner runner = getTestRunner(true);
  MockFlowFile flowFile = new MockFlowFile(1);
  Map<String, String> attributes = new HashMap<>();
  attributes.put("riemann.metric", "5");
  flowFile.putAttributes(attributes);
  runner.enqueue(flowFile);
  try {
    runner.run();
  } catch (ProcessException e) {
    runner.assertQueueNotEmpty();
    throw e;
  }
}
 

@Test
public void testIOException() throws Exception {
    Path filePath = new Path("/some/path/to/file.txt");
    when(mockFileSystem.exists(any(Path.class))).thenThrow(new IOException());
    DeleteHDFS deleteHDFS = new TestableDeleteHDFS(kerberosProperties, mockFileSystem);
    TestRunner runner = TestRunners.newTestRunner(deleteHDFS);
    runner.setProperty(DeleteHDFS.FILE_OR_DIRECTORY, "${hdfs.file}");
    Map<String, String> attributes = Maps.newHashMap();
    attributes.put("hdfs.file", filePath.toString());
    runner.enqueue("foo", attributes);
    runner.run();
    runner.assertQueueNotEmpty();
    runner.assertPenalizeCount(1);
    assertEquals(1, runner.getQueueSize().getObjectCount());
}
 

@Test
public void testStateFailures() throws InterruptedException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(AttributeRollingWindow.class);
    MockStateManager mockStateManager = runner.getStateManager();
    final AttributeRollingWindow processor = (AttributeRollingWindow) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();

    runner.setProperty(AttributeRollingWindow.VALUE_TO_TRACK, "${value}");
    runner.setProperty(AttributeRollingWindow.TIME_WINDOW, "3 sec");

    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("value", "1");

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);

    runner.enqueue(new byte[0],attributes);
    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(AttributeRollingWindow.REL_FAILED_SET_STATE, 1);
    MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(REL_FAILED_SET_STATE).get(0);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_VALUE_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_COUNT_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_MEAN_KEY);
}
 

@Test
public void testStateFailures() throws InterruptedException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(AttributeRollingWindow.class);
    MockStateManager mockStateManager = runner.getStateManager();
    final AttributeRollingWindow processor = (AttributeRollingWindow) runner.getProcessor();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();

    runner.setProperty(AttributeRollingWindow.VALUE_TO_TRACK, "${value}");
    runner.setProperty(AttributeRollingWindow.TIME_WINDOW, "3 sec");

    processor.onScheduled(runner.getProcessContext());

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("value", "1");

    mockStateManager.setFailOnStateGet(Scope.LOCAL, true);

    runner.enqueue(new byte[0],attributes);
    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(AttributeRollingWindow.REL_FAILED_SET_STATE, 1);
    MockFlowFile mockFlowFile = runner.getFlowFilesForRelationship(REL_FAILED_SET_STATE).get(0);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_VALUE_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_COUNT_KEY);
    mockFlowFile.assertAttributeNotExists(ROLLING_WINDOW_MEAN_KEY);
}
 

@Test
public void testBatchLimit() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.FLOWFILE_RATE);
    runner.setProperty(ControlRate.MAX_RATE, "5555");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");

    final int TEST_FILE_COUNT = 1500;

    for (int i = 0; i < TEST_FILE_COUNT; i++) {
        runner.enqueue("test data " + i);
    }

    runner.run(1, false);

    // after 1 run should have MAX_FLOW_FILES_PER_BATCH files transferred and remainder of TEST_FILE_COUNT in queue
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, MAX_FLOW_FILES_PER_BATCH);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();
    assertEquals(TEST_FILE_COUNT - MAX_FLOW_FILES_PER_BATCH, runner.getQueueSize().getObjectCount());

    runner.run(1, false);

    // after 2 runs should have TEST_FILE_COUNT files transferred and 0 in queue
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, TEST_FILE_COUNT);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueEmpty();
}
 

@Test
public void testViaAttribute() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.ATTRIBUTE_RATE);
    runner.setProperty(ControlRate.RATE_CONTROL_ATTRIBUTE_NAME, "count");
    runner.setProperty(ControlRate.MAX_RATE, "20000");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");

    createFlowFile(runner, 1000);
    createFlowFile(runner, 3000);
    createFlowFile(runner, 5000);
    createFlowFile(runner, 20000);
    createFlowFile(runner, 1000);

    runner.run(5, false);

    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 4);
    runner.clearTransferState();

    // at this point, we have sent through 29,000 but our max is 20,000 per second.
    // After 1.45 seconds (29000 / 20000), we should be able to send another 20,000
    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();
    Thread.sleep(1200L);

    // at this point, more than TIME_PERIOD 1.0 seconds but less than 1.45 seconds have passed
    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();
    Thread.sleep(600L);

    // at this point, more than 1.45 seconds have passed, so we should be able to send another 20,000
    runner.run();
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 1);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueEmpty();
}
 

@Test
public void testFileCountWithGrouping() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new ControlRate());
    runner.setProperty(ControlRate.RATE_CONTROL_CRITERIA, ControlRate.FLOWFILE_RATE);
    runner.setProperty(ControlRate.MAX_RATE, "2");
    runner.setProperty(ControlRate.TIME_PERIOD, "1 sec");
    runner.setProperty(ControlRate.GROUPING_ATTRIBUTE_NAME, "group");

    createFlowFileWithGroup(runner, "one");
    createFlowFileWithGroup(runner, "two");
    createFlowFileWithGroup(runner, "one");
    createFlowFileWithGroup(runner, "two");
    createFlowFileWithGroup(runner, "one");
    createFlowFileWithGroup(runner, "two");

    runner.run(6, false);

    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 4);
    runner.clearTransferState();

    runner.run(50, false);
    runner.assertTransferCount(ControlRate.REL_SUCCESS, 0);
    runner.assertTransferCount(ControlRate.REL_FAILURE, 0);
    runner.assertQueueNotEmpty();

    // we have sent 2 files per group and after 1 second, we should be able to send the remaining 1 file per group
    Thread.sleep(1100L);
    runner.run(2);
    runner.assertAllFlowFilesTransferred(ControlRate.REL_SUCCESS, 2);
    runner.assertQueueEmpty();
}
 

@Test
public void testStateFailuresWithRulesUsingOriginal() throws Exception {
    final Criteria criteria = getCriteria();
    criteria.setFlowFilePolicy(FlowFilePolicy.USE_ORIGINAL);
    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).assertAttributeEquals("maxValue2", "1");
}
 

@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");
}
 

@Test
public void testOldestBinIsExpired() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_BIN_AGE, "1 day");
    runner.setProperty(MergeContent.MAX_BIN_COUNT, "50");
    runner.setProperty(MergeContent.MIN_ENTRIES, "10");
    runner.setProperty(MergeContent.MAX_ENTRIES, "10");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_CONCAT);
    runner.setProperty(MergeContent.CORRELATION_ATTRIBUTE_NAME, "correlationId");

    final Map<String, String> attrs = new HashMap<>();
    for (int i = 0; i < 49; i++) {
        attrs.put("correlationId", String.valueOf(i));

        for (int j = 0; j < 5; j++) {
            runner.enqueue(new byte[0], attrs);
        }
    }

    runner.run();

    runner.assertQueueNotEmpty(); // sessions rolled back.
    runner.assertTransferCount(MergeContent.REL_MERGED, 0);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 0);

    attrs.remove("correlationId");

    runner.clearTransferState();

    // Run a single iteration but do not perform the @OnStopped action because
    // we do not want to purge our Bin Manager. This causes some bins to get
    // created. We then enqueue a FlowFile with no correlation id. We do it this
    // way because if we just run a single iteration, then all FlowFiles will be
    // pulled in at once, and we don't know if the first bin to be created will
    // have 5 FlowFiles or 1 FlowFile, since this one that we are about to enqueue
    // will be in a separate bin.
    runner.run(1, false, true);
    runner.enqueue(new byte[0], attrs);
    runner.run();

    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 5);
}
 

@Test
public void testOldestBinIsExpired() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_BIN_AGE, "1 day");
    runner.setProperty(MergeContent.MAX_BIN_COUNT, "50");
    runner.setProperty(MergeContent.MIN_ENTRIES, "10");
    runner.setProperty(MergeContent.MAX_ENTRIES, "10");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_CONCAT);
    runner.setProperty(MergeContent.CORRELATION_ATTRIBUTE_NAME, "correlationId");

    final Map<String, String> attrs = new HashMap<>();
    for (int i = 0; i < 49; i++) {
        attrs.put("correlationId", String.valueOf(i));

        for (int j = 0; j < 5; j++) {
            runner.enqueue(new byte[0], attrs);
        }
    }

    runner.run();

    runner.assertQueueNotEmpty(); // sessions rolled back.
    runner.assertTransferCount(MergeContent.REL_MERGED, 0);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 0);

    attrs.remove("correlationId");

    runner.clearTransferState();

    // Run a single iteration but do not perform the @OnStopped action because
    // we do not want to purge our Bin Manager. This causes some bins to get
    // created. We then enqueue a FlowFile with no correlation id. We do it this
    // way because if we just run a single iteration, then all FlowFiles will be
    // pulled in at once, and we don't know if the first bin to be created will
    // have 5 FlowFiles or 1 FlowFile, since this one that we are about to enqueue
    // will be in a separate bin.
    runner.run(1, false, true);
    runner.enqueue(new byte[0], attrs);
    runner.run();

    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 5);
}
 

@Test
public void testStateFailuresWithRulesUsingOriginal() throws Exception {
    final Criteria criteria = getCriteria();
    criteria.setFlowFilePolicy(FlowFilePolicy.USE_ORIGINAL);
    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).assertAttributeEquals("maxValue2", "1");
}