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

@Test
public void testTransferSmall() throws IOException {
    final TestRunner testRunner = TestRunners.newTestRunner(new SegmentContent());
    testRunner.setProperty(SegmentContent.SIZE, "4 KB");

    testRunner.enqueue(new byte[]{1, 2, 3, 4, 5, 6, 7, 8, 9});
    testRunner.run();

    testRunner.assertTransferCount(SegmentContent.REL_ORIGINAL, 1);
    final MockFlowFile originalFlowFile = testRunner.getFlowFilesForRelationship(SegmentContent.REL_ORIGINAL).get(0);
    originalFlowFile.assertAttributeExists(SegmentContent.FRAGMENT_ID);
    originalFlowFile.assertAttributeEquals(SegmentContent.FRAGMENT_COUNT, "1");

    testRunner.assertTransferCount(SegmentContent.REL_SEGMENTS, 1);
    final MockFlowFile out1 = testRunner.getFlowFilesForRelationship(SegmentContent.REL_SEGMENTS).get(0);
    out1.assertContentEquals(new byte[]{1, 2, 3, 4, 5, 6, 7, 8, 9});
}
 

@Test
public void testRecordSplitBareOutputWithSingleRecords() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitAvro());
    runner.setProperty(SplitAvro.OUTPUT_STRATEGY, SplitAvro.BARE_RECORD_OUTPUT);

    runner.enqueue(users.toByteArray());
    runner.run();

    runner.assertTransferCount(SplitAvro.REL_SPLIT, 100);
    runner.assertTransferCount(SplitAvro.REL_ORIGINAL, 1);
    runner.assertTransferCount(SplitAvro.REL_FAILURE, 0);

    runner.getFlowFilesForRelationship(SplitAvro.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT.key(), "100");
    final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(SplitAvro.REL_SPLIT);

    checkBareRecordsSplitSize(flowFiles, 1, true);
}
 

@Test
public void testAttribute_includeCoreAttributesContent() throws IOException {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToJSON());
    testRunner.setProperty(AttributesToJSON.DESTINATION, AttributesToJSON.DESTINATION_CONTENT);
    testRunner.setProperty(AttributesToJSON.INCLUDE_CORE_ATTRIBUTES, "true");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();

    testRunner.enqueue(ff);
    testRunner.run();

    List<MockFlowFile> flowFilesForRelationship = testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS);

    testRunner.assertTransferCount(AttributesToJSON.REL_FAILURE, 0);
    testRunner.assertTransferCount(AttributesToJSON.REL_SUCCESS, 1);

    MockFlowFile flowFile = flowFilesForRelationship.get(0);

    assertEquals(AttributesToJSON.APPLICATION_JSON, flowFile.getAttribute(CoreAttributes.MIME_TYPE.key()));

    Map<String, String> val = new ObjectMapper().readValue(flowFile.toByteArray(), HashMap.class);
    assertEquals(3, val.size());
    Set<String> coreAttributes = Arrays.stream(CoreAttributes.values()).map(CoreAttributes::key).collect(Collectors.toSet());
    val.keySet().forEach(k -> assertTrue(coreAttributes.contains(k)));
}
 

@Test
public void testOgg() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractMediaMetadata());
    runner.setProperty(ExtractMediaMetadata.METADATA_KEY_FILTER, "");
    runner.setProperty(ExtractMediaMetadata.METADATA_KEY_PREFIX, "ogg.");
    runner.assertValid();

    runner.enqueue(new File("target/test-classes/testVORBIS.ogg").toPath());
    runner.run();

    runner.assertAllFlowFilesTransferred(ExtractMediaMetadata.SUCCESS, 1);
    runner.assertTransferCount(ExtractMediaMetadata.FAILURE, 0);

    final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ExtractMediaMetadata.SUCCESS);
    MockFlowFile flowFile0 = successFiles.get(0);
    flowFile0.assertAttributeExists("filename");
    flowFile0.assertAttributeEquals("filename", "testVORBIS.ogg");
    flowFile0.assertAttributeExists("ogg.Content-Type");
    assertTrue(flowFile0.getAttribute("ogg.Content-Type").startsWith("audio/vorbis"));
    flowFile0.assertAttributeExists("ogg.X-Parsed-By");
    assertTrue(flowFile0.getAttribute("ogg.X-Parsed-By").contains("org.apache.tika.parser.DefaultParser"));
    assertTrue(flowFile0.getAttribute("ogg.X-Parsed-By").contains("org.gagravarr.tika.VorbisParser"));
}
 

@Test
public void testDeleteWithNullValue() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(ConvertJSONToSQL.class);
    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);

    runner.setProperty(ConvertJSONToSQL.CONNECTION_POOL, "dbcp");
    runner.setProperty(ConvertJSONToSQL.TABLE_NAME, "PERSONS");
    runner.setProperty(ConvertJSONToSQL.STATEMENT_TYPE, "DELETE");
    runner.enqueue(Paths.get("src/test/resources/TestConvertJSONToSQL/person-with-null-code.json"));
    runner.run();

    runner.assertTransferCount(ConvertJSONToSQL.REL_ORIGINAL, 1);
    runner.getFlowFilesForRelationship(ConvertJSONToSQL.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT.key(), "1");
    runner.assertTransferCount(ConvertJSONToSQL.REL_SQL, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ConvertJSONToSQL.REL_SQL).get(0);
    out.assertAttributeEquals("sql.args.1.type", String.valueOf(java.sql.Types.INTEGER));
    out.assertAttributeEquals("sql.args.1.value", "1");
    out.assertAttributeEquals("sql.args.2.type", String.valueOf(java.sql.Types.VARCHAR));
    out.assertAttributeEquals("sql.args.2.value", "Mark");
    out.assertAttributeEquals("sql.args.3.type", String.valueOf(java.sql.Types.INTEGER));
    out.assertAttributeNotExists("sql.args.3.value");

    out.assertContentEquals("DELETE FROM PERSONS WHERE ID = ? AND NAME = ? AND CODE = ?");
}
 

@Test
public void testEqualsNotNullStrNotNullOrEmpty() {
    final TestRunner runner = TestRunners.newTestRunner(new ValidateCsv());
    runner.setProperty(ValidateCsv.DELIMITER_CHARACTER, ",");
    runner.setProperty(ValidateCsv.END_OF_LINE_CHARACTER, "\r\n");
    runner.setProperty(ValidateCsv.QUOTE_CHARACTER, "\"");
    runner.setProperty(ValidateCsv.HEADER, "false");

    runner.setProperty(ValidateCsv.SCHEMA, "Equals(), NotNull(), StrNotNullOrEmpty()");

    runner.enqueue("test,test,test\r\ntest,test,test");
    runner.run();
    runner.assertAllFlowFilesTransferred(ValidateCsv.REL_VALID, 1);

    runner.enqueue("test,test,test\r\ntset,test,test");
    runner.run();
    runner.assertTransferCount(ValidateCsv.REL_INVALID, 1);
}
 

@Test
public void testSplitWithTwoLineHeader() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitText());
    runner.setProperty(SplitText.HEADER_LINE_COUNT, "2");
    runner.setProperty(SplitText.LINE_SPLIT_COUNT, "3");

    runner.enqueue(file);
    runner.run();

    runner.assertTransferCount(SplitText.REL_FAILURE, 0);
    runner.assertTransferCount(SplitText.REL_ORIGINAL, 1);
    runner.assertTransferCount(SplitText.REL_SPLITS, 4);

    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(SplitText.REL_SPLITS);
    for (int i = 0; i < splits.size(); i++) {
        final MockFlowFile split = splits.get(i);
        split.assertContentEquals(file.getParent().resolve((i + 1) + ".txt"));
        split.assertAttributeEquals(SplitText.FRAGMENT_INDEX, String.valueOf(i + 1));
    }
}
 

@Test
public void testBmp() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractMediaMetadata());
    runner.setProperty(ExtractMediaMetadata.METADATA_KEY_PREFIX, "bmp.");
    runner.assertValid();

    runner.enqueue(new File("target/test-classes/16color-10x10.bmp").toPath());
    runner.run(2);

    runner.assertAllFlowFilesTransferred(ExtractMediaMetadata.SUCCESS, 1);
    runner.assertTransferCount(ExtractMediaMetadata.FAILURE, 0);

    final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ExtractMediaMetadata.SUCCESS);
    MockFlowFile flowFile0 = successFiles.get(0);
    flowFile0.assertAttributeExists("filename");
    flowFile0.assertAttributeEquals("filename", "16color-10x10.bmp");
    flowFile0.assertAttributeExists("bmp.Content-Type");
    flowFile0.assertAttributeEquals("bmp.Content-Type", "image/bmp");
    flowFile0.assertAttributeExists("bmp.X-Parsed-By");
    assertTrue(flowFile0.getAttribute("bmp.X-Parsed-By").contains("org.apache.tika.parser.DefaultParser"));
    // assertTrue(flowFile0.getAttribute("bmp.X-Parsed-By").contains("org.apache.tika.parser.image.ImageParser"));
    flowFile0.assertAttributeExists("bmp.height");
    flowFile0.assertAttributeEquals("bmp.height", "10");
    flowFile0.assertAttributeExists("bmp.width");
    flowFile0.assertAttributeEquals("bmp.width", "10");
}
 

@Test
public void testEmptyContent() throws IOException {
    TestRunner runner = TestRunners.newTestRunner(ConvertCSVToAvro.class);
    runner.assertNotValid();
    runner.setProperty(ConvertCSVToAvro.SCHEMA, SCHEMA.toString());
    runner.assertValid();

    runner.enqueue(streamFor(""));
    runner.run();

    long converted = runner.getCounterValue("Converted records");
    long errors = runner.getCounterValue("Conversion errors");
    Assert.assertEquals("Should convert 0 rows", 0, converted);
    Assert.assertEquals("Should reject 0 row", 0, errors);

    runner.assertTransferCount("success", 0);
    runner.assertTransferCount("failure", 1);
    runner.assertTransferCount("incompatible", 0);

    MockFlowFile incompatible = runner.getFlowFilesForRelationship("failure").get(0);
    Assert.assertEquals("Should set an error message",
            "No incoming records", incompatible.getAttribute("errors"));
}
 

@Test
public void testSplit_arrayResult_oneValue() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new SplitJson());
    testRunner.setProperty(SplitJson.ARRAY_JSON_PATH_EXPRESSION, "$[0].range[?(@ == 0)]");

    testRunner.enqueue(JSON_SNIPPET);
    testRunner.run();

    testRunner.assertTransferCount(SplitJson.REL_ORIGINAL, 1);
    testRunner.getFlowFilesForRelationship(SplitJson.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT.key(), "1");
    testRunner.assertTransferCount(SplitJson.REL_SPLIT, 1);
    testRunner.getFlowFilesForRelationship(SplitJson.REL_ORIGINAL).get(0).assertContentEquals(JSON_SNIPPET);
    testRunner.getFlowFilesForRelationship(SplitJson.REL_SPLIT).get(0).assertContentEquals("0");
}
 

@Test
public void testAttribute_noIncludeCoreAttributesUserDefined() throws IOException {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToJSON());
    testRunner.setProperty(AttributesToJSON.ATTRIBUTES_LIST, " " + TEST_ATTRIBUTE_KEY + " , " + CoreAttributes.PATH.key() + " ");
    testRunner.setProperty(AttributesToJSON.INCLUDE_CORE_ATTRIBUTES, "false");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();
    ff = session.putAttribute(ff, TEST_ATTRIBUTE_KEY, TEST_ATTRIBUTE_VALUE);
    ff = session.putAttribute(ff, CoreAttributes.PATH.key(), TEST_ATTRIBUTE_VALUE);

    testRunner.enqueue(ff);
    testRunner.run();

    testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS).get(0)
            .assertAttributeExists(AttributesToJSON.JSON_ATTRIBUTE_NAME);
    testRunner.assertTransferCount(AttributesToJSON.REL_SUCCESS, 1);
    testRunner.assertTransferCount(AttributesToJSON.REL_FAILURE, 0);

    String json = testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS)
            .get(0).getAttribute(AttributesToJSON.JSON_ATTRIBUTE_NAME);

    ObjectMapper mapper = new ObjectMapper();
    Map<String, String> val = mapper.readValue(json, HashMap.class);
    assertEquals(TEST_ATTRIBUTE_VALUE, val.get(TEST_ATTRIBUTE_KEY));
    assertEquals(1, val.size());
}
 

@Test
public void testInsertWithMissingColumnWarning() throws InitializationException, ProcessException, SQLException, IOException {
    final TestRunner runner = TestRunners.newTestRunner(ConvertJSONToSQL.class);

    try (final Connection conn = service.getConnection()) {
        try (final Statement stmt = conn.createStatement()) {
            stmt.executeUpdate("CREATE TABLE PERSONS2 (id integer, name varchar(100), code integer, generated_key integer primary key)");
        }
    }

    runner.addControllerService("dbcp", service);
    runner.enableControllerService(service);
    runner.setProperty(ConvertJSONToSQL.CONNECTION_POOL, "dbcp");
    runner.setProperty(ConvertJSONToSQL.TABLE_NAME, "PERSONS2");
    runner.setProperty(ConvertJSONToSQL.STATEMENT_TYPE, "INSERT");
    runner.setProperty(ConvertJSONToSQL.UNMATCHED_COLUMN_BEHAVIOR, ConvertJSONToSQL.WARNING_UNMATCHED_COLUMN);
    runner.enqueue(Paths.get("src/test/resources/TestConvertJSONToSQL/person-1.json"));
    runner.run();

    runner.assertTransferCount(ConvertJSONToSQL.REL_ORIGINAL, 1);
    runner.getFlowFilesForRelationship(ConvertJSONToSQL.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT.key(), "1");
    runner.assertTransferCount(ConvertJSONToSQL.REL_SQL, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ConvertJSONToSQL.REL_SQL).get(0);
    out.assertAttributeEquals("sql.args.1.type", String.valueOf(java.sql.Types.INTEGER));
    out.assertAttributeEquals("sql.args.1.value", "1");
    out.assertAttributeEquals("sql.args.2.type", String.valueOf(java.sql.Types.VARCHAR));
    out.assertAttributeEquals("sql.args.2.value", "Mark");
    out.assertAttributeEquals("sql.args.3.type", String.valueOf(java.sql.Types.INTEGER));
    out.assertAttributeEquals("sql.args.3.value", "48");

    out.assertContentEquals("INSERT INTO PERSONS2 (ID, NAME, CODE) VALUES (?, ?, ?)");
}
 

@Test
public void testValidEmailWithoutAttachments() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractEmailAttachments());

    // Create the message dynamically
    byte [] simpleEmail = attachmentGenerator.SimpleEmail();

    runner.enqueue(simpleEmail);
    runner.run();

    runner.assertTransferCount(ExtractEmailAttachments.REL_ORIGINAL, 1);
    runner.assertTransferCount(ExtractEmailAttachments.REL_FAILURE, 0);
    runner.assertTransferCount(ExtractEmailAttachments.REL_ATTACHMENTS, 0);

}
 

@Test
public void testAclOwnerGroup() throws Exception {
    reset(storage);
    final TestRunner runner = buildNewRunner(getProcessor());
    addRequiredPropertiesToRunner(runner);
    runner.assertValid();

    final Blob blob = mock(Blob.class);

    final Acl.Group mockGroup = mock(Acl.Group.class);
    when(mockGroup.getEmail()).thenReturn(OWNER_GROUP_EMAIL);
    when(blob.getOwner()).thenReturn(mockGroup);

    when(storage.get(any(BlobId.class))).thenReturn(blob);
    when(storage.reader(any(BlobId.class), any(Storage.BlobSourceOption.class))).thenReturn(new MockReadChannel(CONTENT));


    runner.enqueue("");

    runner.run();

    verify(storage).get(any(BlobId.class));
    verify(storage).reader(any(BlobId.class), any(Storage.BlobSourceOption.class));

    runner.assertAllFlowFilesTransferred(FetchGCSObject.REL_SUCCESS);
    runner.assertTransferCount(FetchGCSObject.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(FetchGCSObject.REL_SUCCESS).get(0);

    assertEquals(
            OWNER_GROUP_EMAIL,
            flowFile.getAttribute(OWNER_ATTR)
    );

    assertEquals(
            "group",
            flowFile.getAttribute(OWNER_TYPE_ATTR)
    );
}
 

@Test
public void testSchemaWithCoreAttribuesToAttribute() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToCSV());
    testRunner.setProperty(AttributesToCSV.DESTINATION, OUTPUT_NEW_ATTRIBUTE);
    testRunner.setProperty(AttributesToCSV.INCLUDE_CORE_ATTRIBUTES, "true");
    testRunner.setProperty(AttributesToCSV.NULL_VALUE_FOR_EMPTY_STRING, "false");
    testRunner.setProperty(AttributesToCSV.INCLUDE_SCHEMA, "true");
    testRunner.setProperty(AttributesToCSV.ATTRIBUTES_REGEX, "beach-.*");

    Map<String, String> attrs = new HashMap<String, String>(){{
        put("beach-name", "Malibu Beach");
        put("beach-location", "California, US");
        put("attribute-should-be-eliminated", "This should not be in CSVData!");
    }};

    testRunner.enqueue(new byte[0], attrs);
    testRunner.run();

    testRunner.assertTransferCount(AttributesToCSV.REL_SUCCESS, 1);
    testRunner.assertTransferCount(AttributesToCSV.REL_FAILURE, 0);

    MockFlowFile flowFile = testRunner.getFlowFilesForRelationship(AttributesToCSV.REL_SUCCESS).get(0);
    flowFile.assertAttributeExists("CSVData");
    flowFile.assertAttributeExists("CSVSchema");

    final String path = flowFile.getAttribute("path");
    final String filename = flowFile.getAttribute("filename");
    final String uuid = flowFile.getAttribute("uuid");

    flowFile.assertAttributeEquals("CSVData", "Malibu Beach,\"California, US\"," + path + "," + filename + "," + uuid);
    flowFile.assertAttributeEquals("CSVSchema","beach-name,beach-location,path,filename,uuid");
}
 

@Test
public void testSplit_pathToArrayWithNulls_emptyStringRepresentation() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new SplitJson());
    testRunner.setProperty(SplitJson.ARRAY_JSON_PATH_EXPRESSION, "$.arrayOfNulls");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();

    ff = session.write(ff, new OutputStreamCallback() {
        @Override
        public void process(OutputStream out) throws IOException {
            try (OutputStream outputStream = new BufferedOutputStream(out)) {
                outputStream.write("{\"stringField\": \"String Value\", \"arrayOfNulls\": [null, null, null]}".getBytes(StandardCharsets.UTF_8));
            }
        }
    });

    testRunner.enqueue(ff);
    testRunner.run();

    /* assert that three files were transferred to split and each is empty */
    int expectedFiles = 3;
    testRunner.assertTransferCount(SplitJson.REL_SPLIT, expectedFiles);
    for (int i = 0; i < expectedFiles; i++) {
        testRunner.getFlowFilesForRelationship(SplitJson.REL_SPLIT).get(i).assertContentEquals("");
    }
}
 

@Test
public void testOneLineSplitWithHeader() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitText());
    runner.setProperty(SplitText.HEADER_LINE_COUNT, "1");
    runner.setProperty(SplitText.LINE_SPLIT_COUNT, "1");

    runner.enqueue(TEST_INPUT_DATA);
    runner.run();

    runner.assertTransferCount(SplitText.REL_FAILURE, 0);
    runner.assertTransferCount(SplitText.REL_ORIGINAL, 1);
    runner.assertTransferCount(SplitText.REL_SPLITS, 47);

    final List<MockFlowFile> splits = runner.getFlowFilesForRelationship(SplitText.REL_SPLITS);

    splits.get(0).assertContentEquals("HeaderLine1\nLine2SpacesAtEnd  ");
    splits.get(1).assertContentEquals("HeaderLine1\nLine3");
    splits.get(2).assertContentEquals("HeaderLine1\nLine4");
    splits.get(3).assertContentEquals("HeaderLine1");
    splits.get(4).assertContentEquals("HeaderLine1");
    splits.get(5).assertContentEquals("HeaderLine1\nLine8");
    splits.get(6).assertContentEquals("HeaderLine1\nLine9");
    splits.get(7).assertContentEquals("HeaderLine1");
    splits.get(8).assertContentEquals("HeaderLine1");
    splits.get(9).assertContentEquals("HeaderLine1\n13");
    splits.get(10).assertContentEquals("HeaderLine1\n14");
    splits.get(11).assertContentEquals("HeaderLine1\n15    EndofLine15");
    splits.get(12).assertContentEquals("HeaderLine1\n16");
    for (int i = 13; i < 46; i++) {
        splits.get(i).assertContentEquals("HeaderLine1");
    }
    splits.get(46).assertContentEquals("HeaderLine1\nLastLine");
}
 

@Test
public void test() {
    final TestRunner runner = TestRunners.newTestRunner(new HashAttribute());
    runner.setProperty(HashAttribute.HASH_VALUE_ATTRIBUTE.getName(), "hashValue");
    runner.setProperty("MDKey1", ".*");
    runner.setProperty("MDKey2", "(.).*");

    final Map<String, String> attributeMap = new HashMap<>();
    attributeMap.put("MDKey1", "a");
    attributeMap.put("MDKey2", "b");
    runner.enqueue(new byte[0], attributeMap);

    attributeMap.put("MDKey1", "1");
    attributeMap.put("MDKey2", "2");
    runner.enqueue(new byte[0], attributeMap);

    attributeMap.put("MDKey1", "a");
    attributeMap.put("MDKey2", "z");
    runner.enqueue(new byte[0], attributeMap);

    attributeMap.put("MDKey1", "a");
    attributeMap.put("MDKey2", "bad");
    runner.enqueue(new byte[0], attributeMap);

    attributeMap.put("MDKey1", "a");
    attributeMap.remove("MDKey2");
    runner.enqueue(new byte[0], attributeMap);

    runner.run(5);

    runner.assertTransferCount(HashAttribute.REL_FAILURE, 1);
    runner.assertTransferCount(HashAttribute.REL_SUCCESS, 4);

    final List<MockFlowFile> success = runner.getFlowFilesForRelationship(HashAttribute.REL_SUCCESS);
    final Map<String, Integer> correlationCount = new HashMap<>();
    for (final MockFlowFile flowFile : success) {
        final String correlationId = flowFile.getAttribute("hashValue");
        assertNotNull(correlationId);

        Integer cur = correlationCount.get(correlationId);
        if (cur == null) {
            cur = 0;
        }

        correlationCount.put(correlationId, cur + 1);
    }

    int twoCount = 0;
    int oneCount = 0;
    for (final Integer i : correlationCount.values()) {
        if (i == 1) {
            oneCount++;
        } else if (i == 2) {
            twoCount++;
        } else {
            fail("Got count of " + i);
        }
    }

    assertEquals(1, twoCount);
    assertEquals(2, oneCount);
}
 

@Test
public void testFirstMessage() throws InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new TestableProcessor(1000L));
    runner.setProperty(MonitorActivity.CONTINUALLY_SEND_MESSAGES, "false");
    runner.setProperty(MonitorActivity.THRESHOLD, "100 millis");

    runner.enqueue(new byte[0]);
    runner.run();
    runner.assertAllFlowFilesTransferred(MonitorActivity.REL_SUCCESS, 1);
    runner.clearTransferState();

    Thread.sleep(1000L);

    runNext(runner);
    runner.assertAllFlowFilesTransferred(MonitorActivity.REL_INACTIVE, 1);
    runner.clearTransferState();

    // ensure we don't keep creating the message
    for (int i = 0; i < 10; i++) {
        runNext(runner);
        runner.assertTransferCount(MonitorActivity.REL_SUCCESS, 0);
        runner.assertTransferCount(MonitorActivity.REL_INACTIVE, 0);
        runner.assertTransferCount(MonitorActivity.REL_ACTIVITY_RESTORED, 0);
        Thread.sleep(100L);
    }

    Map<String, String> attributes = new HashMap<>();
    attributes.put("key", "value");
    attributes.put("key1", "value1");

    runner.enqueue(new byte[0], attributes);
    runNext(runner);

    runner.assertTransferCount(MonitorActivity.REL_SUCCESS, 1);
    runner.assertTransferCount(MonitorActivity.REL_ACTIVITY_RESTORED, 1);

    MockFlowFile restoredFlowFile = runner.getFlowFilesForRelationship(MonitorActivity.REL_ACTIVITY_RESTORED).get(0);
    String flowFileContent = new String(restoredFlowFile.toByteArray());
    Assert.assertTrue(Pattern.matches("Activity restored at time: (.*) after being inactive for 0 minutes", flowFileContent));
    restoredFlowFile.assertAttributeNotExists("key");
    restoredFlowFile.assertAttributeNotExists("key1");

    runner.clearTransferState();
    runner.setProperty(MonitorActivity.CONTINUALLY_SEND_MESSAGES, "true");
    Thread.sleep(200L);

    for (int i = 0; i < 10; i++) {
        runNext(runner);
        Thread.sleep(200L);
    }

    runner.assertTransferCount(MonitorActivity.REL_INACTIVE, 10);
    runner.assertTransferCount(MonitorActivity.REL_ACTIVITY_RESTORED, 0);
    runner.assertTransferCount(MonitorActivity.REL_SUCCESS, 0);
    runner.clearTransferState();

    runner.enqueue(new byte[0], attributes);
    runNext(runner);

    runner.assertTransferCount(MonitorActivity.REL_INACTIVE, 0);
    runner.assertTransferCount(MonitorActivity.REL_ACTIVITY_RESTORED, 1);
    runner.assertTransferCount(MonitorActivity.REL_SUCCESS, 1);

    restoredFlowFile = runner.getFlowFilesForRelationship(MonitorActivity.REL_ACTIVITY_RESTORED).get(0);
    flowFileContent = new String(restoredFlowFile.toByteArray());
    Assert.assertTrue(Pattern.matches("Activity restored at time: (.*) after being inactive for 0 minutes", flowFileContent));
    restoredFlowFile.assertAttributeNotExists("key");
    restoredFlowFile.assertAttributeNotExists("key1");
}
 

@Test
public void testSuccessfulPutOperation() throws Exception {
    reset(storage, blob);
    final PutGCSObject processor = getProcessor();
    final TestRunner runner = buildNewRunner(processor);
    addRequiredPropertiesToRunner(runner);

    runner.setProperty(PutGCSObject.KEY, KEY);
    runner.setProperty(PutGCSObject.CONTENT_TYPE, CONTENT_TYPE);
    runner.setProperty(PutGCSObject.MD5, MD5);
    runner.setProperty(PutGCSObject.CRC32C, CRC32C);
    runner.setProperty(PutGCSObject.ACL, ACL.name());
    runner.setProperty(PutGCSObject.ENCRYPTION_KEY, ENCRYPTION_KEY);
    runner.setProperty(PutGCSObject.OVERWRITE, String.valueOf(OVERWRITE));
    runner.setProperty(PutGCSObject.CONTENT_DISPOSITION_TYPE, CONTENT_DISPOSITION_TYPE);

    runner.assertValid();

    when(storage.create(blobInfoArgumentCaptor.capture(),
            inputStreamArgumentCaptor.capture(),
            blobWriteOptionArgumentCaptor.capture())).thenReturn(blob);

    runner.enqueue("test", ImmutableMap.of(CoreAttributes.FILENAME.key(), FILENAME));
    runner.run();

    runner.assertAllFlowFilesTransferred(PutGCSObject.REL_SUCCESS);
    runner.assertTransferCount(PutGCSObject.REL_SUCCESS, 1);

    /*

    String text;
    try (final Reader reader = new InputStreamReader(inputStreamArgumentCaptor.getValue())) {
        text = CharStreams.toString(reader);
    }

    assertEquals(
            "FlowFile content should be equal to the Blob content",
            "test",
            text
    );

    */

    final BlobInfo blobInfo = blobInfoArgumentCaptor.getValue();
    assertEquals(
            BUCKET,
            blobInfo.getBucket()
    );

    assertEquals(
            KEY,
            blobInfo.getName()
    );

    assertEquals(
            CONTENT_DISPOSITION_TYPE + "; filename=" + FILENAME,
            blobInfo.getContentDisposition()
    );

    assertEquals(
            CONTENT_TYPE,
            blobInfo.getContentType()
    );

    assertEquals(
            MD5,
            blobInfo.getMd5()
    );

    assertEquals(
            CRC32C,
            blobInfo.getCrc32c()
    );

    assertNull(blobInfo.getMetadata());

    final List<Storage.BlobWriteOption> blobWriteOptions = blobWriteOptionArgumentCaptor.getAllValues();
    final Set<Storage.BlobWriteOption> blobWriteOptionSet = ImmutableSet.copyOf(blobWriteOptions);

    assertEquals(
            "Each of the BlobWriteOptions should be unique",
            blobWriteOptions.size(),
            blobWriteOptionSet.size()
    );

    assertTrue("The doesNotExist BlobWriteOption should be set if OVERWRITE is false",
            blobWriteOptionSet.contains(Storage.BlobWriteOption.doesNotExist()));
    assertTrue("The md5Match BlobWriteOption should be set if MD5 is non-null",
            blobWriteOptionSet.contains(Storage.BlobWriteOption.md5Match()));
    assertTrue("The crc32cMatch BlobWriteOption should be set if CRC32C is non-null",
            blobWriteOptionSet.contains(Storage.BlobWriteOption.crc32cMatch()));
    assertTrue("The predefinedAcl BlobWriteOption should be set if ACL is non-null",
            blobWriteOptionSet.contains(Storage.BlobWriteOption.predefinedAcl(ACL)));
    assertTrue("The encryptionKey BlobWriteOption should be set if ENCRYPTION_KEY is non-null",
            blobWriteOptionSet.contains(Storage.BlobWriteOption.encryptionKey(ENCRYPTION_KEY)));

}