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

@Test
public void testAttrListNoCoreNullOffTwoNewAttrToAttribute() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToCSV());
    testRunner.setProperty(AttributesToCSV.DESTINATION, OUTPUT_NEW_ATTRIBUTE);
    testRunner.setProperty(AttributesToCSV.INCLUDE_CORE_ATTRIBUTES, "false");
    testRunner.setProperty(AttributesToCSV.NULL_VALUE_FOR_EMPTY_STRING, "false");

    final String NON_PRESENT_ATTRIBUTE_KEY = "beach-type,beach-length";
    testRunner.setProperty(AttributesToCSV.ATTRIBUTES_LIST, NON_PRESENT_ATTRIBUTE_KEY);
    testRunner.enqueue(new byte[0]);
    testRunner.run();

    testRunner.getFlowFilesForRelationship(AttributesToCSV.REL_SUCCESS).get(0)
            .assertAttributeExists("CSVData");
    testRunner.assertTransferCount(AttributesToCSV.REL_SUCCESS, 1);
    testRunner.assertTransferCount(AttributesToCSV.REL_FAILURE, 0);

    testRunner.getFlowFilesForRelationship(AttributesToCSV.REL_SUCCESS)
            .get(0).assertAttributeEquals("CSVData",",");
}
 

@Test
public void testRegexWithGoodCaptureGroup() throws IOException {
    // Test the new Default Regex and with a custom Replacement Values that should succeed.
    final TestRunner runner = getRunner();
    runner.setProperty(ReplaceText.SEARCH_VALUE, "(?s)(^.*$)");
    runner.setProperty(ReplaceText.REPLACEMENT_VALUE, "${'$1':toUpper()}"); // will uppercase group with good Java regex
    runner.setProperty(ReplaceText.REPLACEMENT_STRATEGY, ReplaceText.REGEX_REPLACE);
    runner.setProperty(ReplaceText.EVALUATION_MODE, ReplaceText.ENTIRE_TEXT);

    runner.enqueue("testing\n123".getBytes());
    runner.run();

    runner.assertAllFlowFilesTransferred(ReplaceText.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ReplaceText.REL_SUCCESS).get(0);
    out.assertContentEquals("TESTING\n123");
}
 

@Test
public void testUpdateBySimpleKey() {
    TestRunner runner = init(PutMongo.class);
    Document document = new Document()
        .append("name", "John Smith")
        .append("department", "Engineering");
    collection.insertOne(document);

    String updateBody = "{\n" +
        "\t\"name\": \"John Smith\",\n" +
        "\t\"$set\": {\n" +
        "\t\t\"email\": \"[email protected]\",\n" +
        "\t\t\"grade\": \"Sr. Principle Eng.\"\n" +
        "\t},\n" +
        "\t\"$inc\": {\n" +
        "\t\t\"writes\": 1\n" +
        "\t}\n" +
        "}";
    runner.setProperty(PutMongo.UPDATE_QUERY_KEY, "name");
    runner.setProperty(PutMongo.UPDATE_MODE, PutMongo.UPDATE_WITH_OPERATORS);
    runner.setProperty(PutMongo.MODE, PutMongo.MODE_UPDATE);
    runner.setValidateExpressionUsage(true);
    runner.enqueue(updateBody);
    updateTests(runner, document);
}
 

@Test
public void testTransformInputWithDefaultrExpressionLanguage() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new JoltTransformJSON());
    runner.setValidateExpressionUsage(false);
    final String spec = new String(Files.readAllBytes(Paths.get("src/test/resources/TestJoltTransformJson/defaultrELSpec.json")));
    runner.setProperty(JoltTransformJSON.JOLT_SPEC, spec);
    runner.setProperty(JoltTransformJSON.JOLT_TRANSFORM, JoltTransformJSON.DEFAULTR);
    runner.setVariable("quota","5");
    runner.enqueue(JSON_INPUT);
    runner.run();
    runner.assertAllFlowFilesTransferred(JoltTransformJSON.REL_SUCCESS);
    final MockFlowFile transformed = runner.getFlowFilesForRelationship(JoltTransformJSON.REL_SUCCESS).get(0);
    Object transformedJson = JsonUtils.jsonToObject(new ByteArrayInputStream(transformed.toByteArray()));
    Object compareJson = JsonUtils.jsonToObject(Files.newInputStream(Paths.get("src/test/resources/TestJoltTransformJson/defaultrELOutput.json")));
    assertTrue(DIFFY.diff(compareJson, transformedJson).isEmpty());
}
 

@Test
public void testMultipleAvroMessages() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ConvertAvroToJSON());
    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/user.avsc"));

    runner.setProperty(ConvertAvroToJSON.CONTAINER_OPTIONS, ConvertAvroToJSON.CONTAINER_ARRAY);

    final GenericRecord user1 = new GenericData.Record(schema);
    user1.put("name", "Alyssa");
    user1.put("favorite_number", 256);

    final GenericRecord user2 = new GenericData.Record(schema);
    user2.put("name", "George");
    user2.put("favorite_number", 1024);
    user2.put("favorite_color", "red");

    final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
    final ByteArrayOutputStream out1 = AvroTestUtil.serializeAvroRecord(schema, datumWriter, user1, user2);
    runner.enqueue(out1.toByteArray());

    runner.run();

    runner.assertAllFlowFilesTransferred(ConvertAvroToJSON.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ConvertAvroToJSON.REL_SUCCESS).get(0);
    out.assertContentEquals("[{\"name\": \"Alyssa\", \"favorite_number\": 256, \"favorite_color\": null},{\"name\": \"George\", \"favorite_number\": 1024, \"favorite_color\": \"red\"}]");
}
 

@Test
public void testExecuteIngestAndUpdateWithWorkingDirPutToAttribute() throws IOException {
    File exJar = new File("src/test/resources/ExecuteCommand/TestIngestAndUpdate.jar");
    File dummy = new File("src/test/resources/ExecuteCommand/1000bytes.txt");
    String jarPath = exJar.getAbsolutePath();
    exJar.setExecutable(true);
    final TestRunner controller = TestRunners.newTestRunner(ExecuteStreamCommand.class);
    controller.enqueue(dummy.toPath());
    controller.setProperty(ExecuteStreamCommand.WORKING_DIR, "target");
    controller.setProperty(ExecuteStreamCommand.EXECUTION_COMMAND, "java");
    controller.setProperty(ExecuteStreamCommand.PUT_OUTPUT_IN_ATTRIBUTE, "streamOutput");
    controller.setProperty(ExecuteStreamCommand.EXECUTION_ARGUMENTS, "-jar;" + jarPath);
    controller.run(1);
    controller.assertTransferCount(ExecuteStreamCommand.ORIGINAL_RELATIONSHIP, 1);
    List<MockFlowFile> flowFiles = controller.getFlowFilesForRelationship(ExecuteStreamCommand.ORIGINAL_RELATIONSHIP);
    String result = flowFiles.get(0).getAttribute("streamOutput");

    final String quotedSeparator = Pattern.quote(File.separator);
    assertTrue(Pattern.compile(quotedSeparator + "nifi-standard-processors" + quotedSeparator + "target:ModifiedResult\r?\n").matcher(result).find());
}
 

@Test
public void testSimpleDefaultContains() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new RouteText());
    runner.setProperty(RouteText.MATCH_STRATEGY, RouteText.CONTAINS);
    runner.setProperty("simple", "middle");

    runner.enqueue("start middle end\nnot match".getBytes("UTF-8"));
    runner.run();

    runner.assertTransferCount("simple", 1);
    runner.assertTransferCount("unmatched", 1);
    runner.assertTransferCount("original", 1);
    final MockFlowFile outMatched = runner.getFlowFilesForRelationship("simple").get(0);
    outMatched.assertContentEquals("start middle end\n".getBytes("UTF-8"));
    final MockFlowFile outUnmatched = runner.getFlowFilesForRelationship("unmatched").get(0);
    outUnmatched.assertContentEquals("not match".getBytes("UTF-8"));
}
 

@Test
public void testAttributeToContent() throws IOException {
    final TestRunner runner = getRunner();
    runner.setProperty(ReplaceText.SEARCH_VALUE, ".*");
    runner.setProperty(ReplaceText.REPLACEMENT_VALUE, "${abc}");

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "Good");
    runner.enqueue(Paths.get("src/test/resources/hello.txt"), attributes);

    runner.run();

    runner.assertAllFlowFilesTransferred(ReplaceText.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ReplaceText.REL_SUCCESS).get(0);
    out.assertContentEquals("Good");
}
 

@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 testAppendEndlineCR() throws IOException {
    final TestRunner runner = getRunner();
    runner.setProperty(ReplaceText.EVALUATION_MODE, ReplaceText.LINE_BY_LINE);
    runner.setProperty(ReplaceText.REPLACEMENT_VALUE, "TEST");
    runner.setProperty(ReplaceText.REPLACEMENT_STRATEGY, ReplaceText.APPEND);

    runner.enqueue("Hello \rWorld \r".getBytes());
    runner.run();

    runner.assertAllFlowFilesTransferred(ReplaceText.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ReplaceText.REL_SUCCESS).get(0);
    out.assertContentEquals("Hello TEST\rWorld TEST\r".getBytes("UTF-8"));
}
 

@Test
public void validateFailedPublishAndTransferToFailure() throws Exception {
    PublishAMQP pubProc = new LocalPublishAMQP();
    TestRunner runner = TestRunners.newTestRunner(pubProc);
    runner.setProperty(PublishAMQP.HOST, "injvm");
    runner.setProperty(PublishAMQP.EXCHANGE, "badToTheBone");
    runner.setProperty(PublishAMQP.ROUTING_KEY, "key1");

    runner.enqueue("Hello Joe".getBytes());

    runner.run();

    assertTrue(runner.getFlowFilesForRelationship(PublishAMQP.REL_SUCCESS).isEmpty());
    assertNotNull(runner.getFlowFilesForRelationship(PublishAMQP.REL_FAILURE).get(0));
}
 

@Test
public void basicFileFetch() throws IOException {
    FileSystem results = fakeFtpServer.getFileSystem();

    FileEntry sampleFile = new FileEntry("c:\\data\\randombytes-2");
    sampleFile.setContents("Just some random test test test chocolate");
    results.add(sampleFile);

    // Check file exists
    Assert.assertTrue(results.exists("c:\\data\\randombytes-2"));

    TestRunner runner = TestRunners.newTestRunner(FetchFTP.class);
    runner.setProperty(FetchFTP.HOSTNAME, "${host}");
    runner.setProperty(FetchFTP.USERNAME, "${username}");
    runner.setProperty(FTPTransfer.PASSWORD, password);
    runner.setProperty(FTPTransfer.PORT, "${port}");
    runner.setProperty(FetchFTP.REMOTE_FILENAME, "c:\\data\\randombytes-2");
    runner.setProperty(FetchFTP.COMPLETION_STRATEGY, FetchFTP.COMPLETION_MOVE);
    runner.setProperty(FetchFTP.MOVE_DESTINATION_DIR, "data");


    Map<String, String> attrs = new HashMap<String, String>();
    attrs.put("host", "localhost");
    attrs.put("username", username);
    attrs.put("port", Integer.toString(ftpPort));
    runner.enqueue("", attrs);

    runner.run();

    final MockFlowFile retrievedFile = runner.getFlowFilesForRelationship(FetchFTP.REL_SUCCESS).get(0);
    retrievedFile.assertContentEquals("Just some random test test test chocolate");
}
 

@Test
public void testLargePutToAttribute() throws IOException {
    File dummy = new File("src/test/resources/ExecuteCommand/1000bytes.txt");
    File dummy10MBytes = new File("target/10MB.txt");
    byte[] bytes = Files.readAllBytes(dummy.toPath());
    try (FileOutputStream fos = new FileOutputStream(dummy10MBytes)) {
        for (int i = 0; i < 10000; i++) {
            fos.write(bytes, 0, 1000);
        }
    }

    final TestRunner controller = TestRunners.newTestRunner(ExecuteStreamCommand.class);
    controller.setValidateExpressionUsage(false);
    controller.enqueue("".getBytes());
    if(isWindows()) {
        controller.setProperty(ExecuteStreamCommand.EXECUTION_COMMAND, "cmd.exe");
        controller.setProperty(ExecuteStreamCommand.EXECUTION_ARGUMENTS, "/c;type " + dummy10MBytes.getAbsolutePath());
        controller.setProperty(ExecuteStreamCommand.ARG_DELIMITER, ";");
    } else{
        controller.setProperty(ExecuteStreamCommand.EXECUTION_COMMAND, "cat");
        controller.setProperty(ExecuteStreamCommand.EXECUTION_ARGUMENTS, dummy10MBytes.getAbsolutePath());
    }
    controller.setProperty(ExecuteStreamCommand.IGNORE_STDIN, "true");
    controller.setProperty(ExecuteStreamCommand.PUT_OUTPUT_IN_ATTRIBUTE, "executeStreamCommand.output");
    controller.setProperty(ExecuteStreamCommand.PUT_ATTRIBUTE_MAX_LENGTH, "256");

    controller.run(1);
    controller.assertTransferCount(ExecuteStreamCommand.ORIGINAL_RELATIONSHIP, 1);
    controller.assertTransferCount(ExecuteStreamCommand.OUTPUT_STREAM_RELATIONSHIP, 0);
    List<MockFlowFile> flowFiles = controller.getFlowFilesForRelationship(ExecuteStreamCommand.ORIGINAL_RELATIONSHIP);

    flowFiles.get(0).assertAttributeEquals("execution.status", "0");
    String result = flowFiles.get(0).getAttribute("executeStreamCommand.output");
    assertTrue(Pattern.compile("a{256}").matcher(result).matches());
}
 

@Test
public void testPutS3ObjectUsingCredentialsProviderService() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(new PutS3Object());

    final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();

    runner.addControllerService("awsCredentialsProvider", serviceImpl);

    runner.setProperty(serviceImpl, AbstractAWSCredentialsProviderProcessor.CREDENTIALS_FILE, System.getProperty("user.home") + "/aws-credentials.properties");
    runner.enableControllerService(serviceImpl);

    runner.assertValid(serviceImpl);

    runner.setProperty(PutS3Object.AWS_CREDENTIALS_PROVIDER_SERVICE, "awsCredentialsProvider");
    runner.setProperty(PutS3Object.REGION, REGION);
    runner.setProperty(PutS3Object.BUCKET, BUCKET_NAME);

    Assert.assertTrue(runner.setProperty("x-custom-prop", "hello").isValid());

    for (int i = 0; i < 3; i++) {
        final Map<String, String> attrs = new HashMap<>();
        attrs.put("filename", String.valueOf(i) + ".txt");
        runner.enqueue(getResourcePath(SAMPLE_FILE_RESOURCE_NAME), attrs);
    }
    runner.run(3);

    runner.assertAllFlowFilesTransferred(PutS3Object.REL_SUCCESS, 3);

}
 

@Test
public void testSuccessfulParseToAttributes() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ParseCEF());
    runner.setProperty(ParseCEF.FIELDS_DESTINATION, ParseCEF.DESTINATION_ATTRIBUTES);
    runner.enqueue(sample1.getBytes());
    runner.run();

    runner.assertAllFlowFilesTransferred(ParseCEF.REL_SUCCESS, 1);
    final MockFlowFile mff = runner.getFlowFilesForRelationship(ParseCEF.REL_SUCCESS).get(0);
    mff.assertAttributeEquals("cef.extension.rt", sdf.format(new Date(1423441663000L)));
    mff.assertAttributeEquals("cef.extension.cn3Label", "Test Long");
    mff.assertAttributeEquals("cef.extension.cn3", "9223372036854775807");
    mff.assertAttributeEquals("cef.extension.cfp1", "1.234");
    mff.assertAttributeEquals("cef.extension.cfp1Label", "Test FP Number");
    mff.assertAttributeEquals("cef.extension.smac", "00:00:0c:07:ac:00");
    mff.assertAttributeEquals("cef.extension.c6a3", "2001:cdba:0:0:0:0:3257:9652");
    mff.assertAttributeEquals("cef.extension.c6a3Label", "Test IPv6");
    mff.assertAttributeEquals("cef.extension.cs1Label", "Test String");
    mff.assertAttributeEquals("cef.extension.cs1", "test test test chocolate");
    mff.assertAttributeEquals("cef.extension.destinationTranslatedAddress", "123.123.123.123");
    mff.assertContentEquals(sample1.getBytes());


    // Converting a field without timezone will always result on render time being dependent
    // on locale of the machine running this test.
    long eventTime = 1423229263000L;
    int offset = TimeZone.getDefault().getOffset(eventTime);
    sdf.setTimeZone(TimeZone.getDefault());

    String prettyEvent = sdf.format(new Date(eventTime - offset));

    mff.assertAttributeEquals("cef.extension.deviceCustomDate1",prettyEvent);
    mff.assertAttributeEquals("cef.extension.dpt", "1234");
    mff.assertAttributeEquals("cef.extension.agt", "123.123.0.124");
    mff.assertAttributeEquals("cef.extension.dlat", "40.366633");
}
 

@Test
public void testPutFileWithException() throws IOException {
    // Refer to comment in the BeforeClass method for an explanation
    assumeTrue(isNotWindows());

    String dirName = "target/testPutFileWrongPermissions";
    File file = new File(dirName);
    file.mkdirs();
    Configuration config = new Configuration();
    FileSystem fs = FileSystem.get(config);
    Path p = new Path(dirName).makeQualified(fs.getUri(), fs.getWorkingDirectory());

    final KerberosProperties testKerberosProperties = kerberosProperties;
    TestRunner runner = TestRunners.newTestRunner(new PutHDFS() {
        @Override
        protected void changeOwner(ProcessContext context, FileSystem hdfs, Path name) {
            throw new ProcessException("Forcing Exception to get thrown in order to verify proper handling");
        }

        @Override
        protected KerberosProperties getKerberosProperties(File kerberosConfigFile) {
            return testKerberosProperties;
        }
    });
    runner.setProperty(PutHDFS.DIRECTORY, dirName);
    runner.setProperty(PutHDFS.CONFLICT_RESOLUTION, "replace");

    try (FileInputStream fis = new FileInputStream("src/test/resources/testdata/randombytes-1");) {
        Map<String, String> attributes = new HashMap<String, String>();
        attributes.put(CoreAttributes.FILENAME.key(), "randombytes-1");
        runner.enqueue(fis, attributes);
        runner.run();
    }

    List<MockFlowFile> failedFlowFiles = runner
            .getFlowFilesForRelationship(new Relationship.Builder().name("failure").build());
    assertFalse(failedFlowFiles.isEmpty());
    assertTrue(failedFlowFiles.get(0).isPenalized());

    fs.delete(p, true);
}
 

@Test
public void testPutProvenanceSendEventTransitUri() throws JMSException {
    final PutJMS putJMS = spy(new PutJMS());
    final TestRunner runnerPut = TestRunners.newTestRunner(putJMS);
    runnerPut.setProperty(JmsProperties.JMS_PROVIDER, TEST_PROVIDER);
    runnerPut.setProperty(JmsProperties.URL, TEST_URL);
    runnerPut.setProperty(JmsProperties.DESTINATION_TYPE, TEST_DEST_TYPE);
    runnerPut.setProperty(JmsProperties.DESTINATION_NAME, TEST_DEST_NAME + testQueueSuffix());
    runnerPut.setProperty(JmsProperties.ATTRIBUTES_TO_JMS_PROPS, "true");

    runnerPut.enqueue("putGetMessage".getBytes());

    runnerPut.run();

    assertEquals(0, runnerPut.getFlowFilesForRelationship(PutJMS.REL_FAILURE).size());
    assertEquals(1, runnerPut.getFlowFilesForRelationship(PutJMS.REL_SUCCESS).size());

    final List<ProvenanceEventRecord> putProvenanceEvents = runnerPut.getProvenanceEvents();

    assertEquals(1, putProvenanceEvents.size());
    // Verify the transitUri is the same as that configured in the properties
    assertEquals(TEST_URL, putProvenanceEvents.get(0).getTransitUri());

    final GetJMSQueue getJmsQueue = new GetJMSQueue();
    final TestRunner runnerGet = TestRunners.newTestRunner(getJmsQueue);
    runnerGet.setProperty(JmsProperties.JMS_PROVIDER, TEST_PROVIDER);
    runnerGet.setProperty(JmsProperties.URL, TEST_URL);
    runnerGet.setProperty(JmsProperties.DESTINATION_NAME, TEST_DEST_NAME + testQueueSuffix());
    runnerGet.setProperty(JmsProperties.ACKNOWLEDGEMENT_MODE, TEST_ACK_MODE);

    runnerGet.run();

    assertEquals(1, runnerGet.getFlowFilesForRelationship(GetJMSQueue.REL_SUCCESS).size());

    final List<MockFlowFile> flowFilesGet = runnerGet.getFlowFilesForRelationship(GetJMSQueue.REL_SUCCESS);

    assertEquals(1, flowFilesGet.size());
    final MockFlowFile successFlowFile = flowFilesGet.get(0);

    successFlowFile.assertContentEquals("putGetMessage");
}
 

@Test
public void testAvroConcatWithDifferentMetadataIgnore() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_ENTRIES, "3");
    runner.setProperty(MergeContent.MIN_ENTRIES, "3");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_AVRO);
    runner.setProperty(MergeContent.METADATA_STRATEGY, MergeContent.METADATA_STRATEGY_IGNORE);

    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/TestMergeContent/user.avsc"));

    final GenericRecord user1 = new GenericData.Record(schema);
    user1.put("name", "Alyssa");
    user1.put("favorite_number", 256);
    final Map<String, String> userMeta1 = new HashMap<String, String>() {{
        put("test_metadata1", "Test 1");
    }};

    final GenericRecord user2 = new GenericData.Record(schema);
    user2.put("name", "Ben");
    user2.put("favorite_number", 7);
    user2.put("favorite_color", "red");
    final Map<String, String> userMeta2 = new HashMap<String, String>() {{
        put("test_metadata1", "Test 2"); // Test non-matching values
    }};

    final GenericRecord user3 = new GenericData.Record(schema);
    user3.put("name", "John");
    user3.put("favorite_number", 5);
    user3.put("favorite_color", "blue");
    final Map<String, String> userMeta3 = new HashMap<String, String>() {{
        put("test_metadata1", "Test 1");
        put("test_metadata2", "Test"); // Test unique
    }};

    final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
    final ByteArrayOutputStream out1 = serializeAvroRecord(schema, user1, datumWriter, userMeta1);
    final ByteArrayOutputStream out2 = serializeAvroRecord(schema, user2, datumWriter, userMeta2);
    final ByteArrayOutputStream out3 = serializeAvroRecord(schema, user3, datumWriter, userMeta3);

    runner.enqueue(out1.toByteArray());
    runner.enqueue(out2.toByteArray());
    runner.enqueue(out3.toByteArray());

    runner.run();
    runner.assertQueueEmpty();
    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 0);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 3);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0);
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/avro-binary");

    // create a reader for the merged content
    byte[] data = runner.getContentAsByteArray(bundle);
    final Map<String, GenericRecord> users = getGenericRecordMap(data, schema, "name");

    Assert.assertEquals(3, users.size());
    Assert.assertTrue(users.containsKey("Alyssa"));
    Assert.assertTrue(users.containsKey("Ben"));
    Assert.assertTrue(users.containsKey("John"));
}
 

@Test
public void testCollectionFunctionsWithCastChoiceWithInts() throws InitializationException {
    final Record record = createHierarchicalArrayRecord();
    record.setValue("height", 30);

    final ArrayListRecordReader recordReader = new ArrayListRecordReader(record.getSchema());
    recordReader.addRecord(record);
    recordReader.addRecord(record);

    final ArrayListRecordWriter writer = new ArrayListRecordWriter(record.getSchema());

    TestRunner runner = getRunner();
    runner.addControllerService("reader", recordReader);
    runner.enableControllerService(recordReader);
    runner.addControllerService("writer", writer);
    runner.enableControllerService(writer);

    runner.setProperty(QueryRecord.RECORD_READER_FACTORY, "reader");
    runner.setProperty(QueryRecord.RECORD_WRITER_FACTORY, "writer");
    runner.setProperty(REL_NAME,
        "SELECT title, name, " +
            "sum(CAST(height AS INT)) as height_total_int, " +
            "sum(CAST(height AS BIGINT)) as height_total_long " +
            "FROM FLOWFILE " +
            "GROUP BY title, name");

    runner.enqueue(new byte[0]);

    runner.run();

    runner.assertTransferCount(REL_NAME, 1);

    final List<Record> written = writer.getRecordsWritten();
    assertEquals(1, written.size());

    final Number height = 60;
    final Record output = written.get(0);
    assertEquals("John Doe", output.getValue("name"));
    assertEquals("Software Engineer", output.getValue("title"));
    assertEquals(height.longValue(), output.getValue("height_total_long"));
    assertEquals(height.intValue(), output.getValue("height_total_int"));
}
 

@Test
public void testAvroConcatWithDifferentSchemas() throws IOException, InterruptedException {
    final TestRunner runner = TestRunners.newTestRunner(new MergeContent());
    runner.setProperty(MergeContent.MAX_ENTRIES, "3");
    runner.setProperty(MergeContent.MIN_ENTRIES, "3");
    runner.setProperty(MergeContent.MERGE_FORMAT, MergeContent.MERGE_FORMAT_AVRO);

    final Schema schema1 = new Schema.Parser().parse(new File("src/test/resources/TestMergeContent/user.avsc"));
    final Schema schema2 = new Schema.Parser().parse(new File("src/test/resources/TestMergeContent/place.avsc"));

    final GenericRecord record1 = new GenericData.Record(schema1);
    record1.put("name", "Alyssa");
    record1.put("favorite_number", 256);

    final GenericRecord record2 = new GenericData.Record(schema2);
    record2.put("name", "Some Place");

    final GenericRecord record3 = new GenericData.Record(schema1);
    record3.put("name", "John");
    record3.put("favorite_number", 5);
    record3.put("favorite_color", "blue");

    final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema1);
    final ByteArrayOutputStream out1 = serializeAvroRecord(schema1, record1, datumWriter);
    final ByteArrayOutputStream out2 = serializeAvroRecord(schema2, record2, datumWriter);
    final ByteArrayOutputStream out3 = serializeAvroRecord(schema1, record3, datumWriter);

    runner.enqueue(out1.toByteArray());
    runner.enqueue(out2.toByteArray());
    runner.enqueue(out3.toByteArray());

    runner.run();
    runner.assertQueueEmpty();
    runner.assertTransferCount(MergeContent.REL_MERGED, 1);
    runner.assertTransferCount(MergeContent.REL_FAILURE, 1);
    runner.assertTransferCount(MergeContent.REL_ORIGINAL, 3);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(MergeContent.REL_MERGED).get(0);
    bundle.assertAttributeEquals(CoreAttributes.MIME_TYPE.key(), "application/avro-binary");

    final byte[] data = runner.getContentAsByteArray(bundle);
    final Map<String, GenericRecord> users = getGenericRecordMap(data, schema1, "name");
    Assert.assertEquals(2, users.size());
    Assert.assertTrue(users.containsKey("Alyssa"));
    Assert.assertTrue(users.containsKey("John"));

    final MockFlowFile failure = runner.getFlowFilesForRelationship(MergeContent.REL_FAILURE).get(0);
    final byte[] failureData = runner.getContentAsByteArray(failure);
    final Map<String, GenericRecord> places = getGenericRecordMap(failureData, schema2, "name");
    Assert.assertEquals(1, places.size());
    Assert.assertTrue(places.containsKey("Some Place"));
}