Java Code Examples for org.apache.nifi.util.MockFlowFile#toByteArray()

@Test
public void testRecordSplitDatafileOutputWithoutMetadata() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitAvro());
    runner.setProperty(SplitAvro.TRANSFER_METADATA, "false");

    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);
    checkDataFileSplitSize(flowFiles, 1, false);

    for (final MockFlowFile flowFile : flowFiles) {
        try (final ByteArrayInputStream in = new ByteArrayInputStream(flowFile.toByteArray());
             final DataFileStream<GenericRecord> reader = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>())) {
            Assert.assertFalse(reader.getMetaKeys().contains(META_KEY1));
            Assert.assertFalse(reader.getMetaKeys().contains(META_KEY2));
            Assert.assertFalse(reader.getMetaKeys().contains(META_KEY3));
        }
    }
}
 

private void checkDataFileTotalSize(List<MockFlowFile> flowFiles, int expectedTotalRecords) throws IOException {
    int count = 0;
    for (final MockFlowFile flowFile : flowFiles) {
        try (final ByteArrayInputStream in = new ByteArrayInputStream(flowFile.toByteArray());
             final DataFileStream<GenericRecord> reader = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>())) {

            GenericRecord record = null;
            while (reader.hasNext()) {
                record = reader.next(record);
                Assert.assertNotNull(record.get("name"));
                Assert.assertNotNull(record.get("favorite_number"));
                count++;
            }
        }
    }
    assertEquals(expectedTotalRecords, count);
}
 

@Test
public void testWithNoMatch() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ReplaceTextWithMapping());
    runner.setProperty(ReplaceTextWithMapping.REGEX, "-(.*?)-");
    runner.setProperty(ReplaceTextWithMapping.MATCHING_GROUP_FOR_LOOKUP_KEY, "1");
    runner.setProperty(ReplaceTextWithMapping.MAPPING_FILE, Paths.get("src/test/resources/TestReplaceTextWithMapping/color-fruit-no-match-mapping.txt").toFile().getAbsolutePath());

    final Path path = Paths.get("src/test/resources/TestReplaceTextWithMapping/colors.txt");
    runner.enqueue(path);
    runner.run();

    runner.assertAllFlowFilesTransferred(ReplaceTextWithMapping.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ReplaceTextWithMapping.REL_SUCCESS).get(0);
    String outputString = new String(out.toByteArray());
    String expected = new String(Files.readAllBytes(path));
    assertEquals(expected, outputString);
}
 

@Test
public void testExpressionLanguageInText3() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ReplaceTextWithMapping());
    final String mappingFile = Paths.get("src/test/resources/TestReplaceTextWithMapping/color-fruit-mapping.txt").toFile().getAbsolutePath();
    runner.setProperty(ReplaceTextWithMapping.MAPPING_FILE, mappingFile);
    runner.setProperty(ReplaceTextWithMapping.REGEX, ".*\\|(.*?)\\|.*");
    runner.setProperty(ReplaceTextWithMapping.MATCHING_GROUP_FOR_LOOKUP_KEY, "1");

    String text = "${foo}|red|${baz}";

    runner.enqueue(text.getBytes());
    runner.run();

    runner.assertAllFlowFilesTransferred(ReplaceTextWithMapping.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ReplaceTextWithMapping.REL_SUCCESS).get(0);
    String outputString = new String(out.toByteArray());
    String expected = "${foo}|apple|${baz}";
    assertEquals(expected, outputString);
}
 

@Test
public void testHead() throws Exception {
    addHandler(new GetOrHeadHandler());

    runner.setProperty(InvokeHTTP.PROP_METHOD, "HEAD");
    runner.setProperty(InvokeHTTP.PROP_URL, url + "/status/200");

    createFlowFiles(runner);

    runner.run();
    runner.assertTransferCount(InvokeHTTP.REL_SUCCESS_REQ, 1);
    runner.assertTransferCount(InvokeHTTP.REL_RESPONSE, 1);
    runner.assertTransferCount(InvokeHTTP.REL_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
    runner.assertPenalizeCount(0);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(InvokeHTTP.REL_SUCCESS_REQ).get(0);
    bundle.assertContentEquals("Hello".getBytes("UTF-8"));
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
    bundle.assertAttributeEquals("Foo", "Bar");

    final MockFlowFile bundle1 = runner.getFlowFilesForRelationship(InvokeHTTP.REL_RESPONSE).get(0);
    bundle1.assertContentEquals("".getBytes("UTF-8"));
    bundle1.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
    bundle1.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
    bundle1.assertAttributeEquals("Foo", "Bar");
    final String actual1 = new String(bundle1.toByteArray(), StandardCharsets.UTF_8);
    final String expected1 = "";
    Assert.assertEquals(expected1, actual1);
}
 

@Test
public void testDelete() throws Exception {
    addHandler(new DeleteHandler());

    runner.setProperty(InvokeHTTP.PROP_METHOD, "DELETE");
    runner.setProperty(InvokeHTTP.PROP_URL, url + "/status/200");

    createFlowFiles(runner);

    runner.run();
    runner.assertTransferCount(InvokeHTTP.REL_SUCCESS_REQ, 1);
    runner.assertTransferCount(InvokeHTTP.REL_RESPONSE, 1);
    runner.assertTransferCount(InvokeHTTP.REL_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
    runner.assertPenalizeCount(0);

    final MockFlowFile bundle = runner.getFlowFilesForRelationship(InvokeHTTP.REL_SUCCESS_REQ).get(0);
    bundle.assertContentEquals("Hello".getBytes("UTF-8"));
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
    bundle.assertAttributeEquals("Foo", "Bar");

    final MockFlowFile bundle1 = runner.getFlowFilesForRelationship(InvokeHTTP.REL_RESPONSE).get(0);
    bundle1.assertContentEquals("".getBytes("UTF-8"));
    bundle1.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
    bundle1.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
    bundle1.assertAttributeEquals("Foo", "Bar");
    final String actual1 = new String(bundle1.toByteArray(), StandardCharsets.UTF_8);
    final String expected1 = "";
    Assert.assertEquals(expected1, actual1);
}
 

/**
 * Demonstrates transforming the JSON object in an incoming FlowFile to output
 * @throws Exception
 */
@Test
public void testTransformGroovy() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new ExecuteScript());
    runner.setValidateExpressionUsage(false);
    runner.setProperty(SCRIPT_ENGINE, "Groovy");
    runner.setProperty(ScriptingComponentUtils.SCRIPT_FILE, "src/test/resources/executescript/content/transform.groovy");
    runner.setProperty(ScriptingComponentUtils.MODULES, "src/test/resources/executescript");
    runner.assertValid();

    InputObject inputJsonObject = new InputObject();
    inputJsonObject.value = 3;
    String inputJson = gson.toJson(inputJsonObject);

    runner.enqueue(inputJson);
    runner.run();

    runner.assertAllFlowFilesTransferred("success", 1);
    final List<MockFlowFile> successFlowFiles = runner.getFlowFilesForRelationship("success");
    MockFlowFile result = successFlowFiles.get(0);
    byte[] flowFileBytes = result.toByteArray();
    String outputJson = new String(flowFileBytes);

    OutputObject outputJsonObject = gson.fromJson(outputJson, OutputObject.class);
    Assert.assertEquals(9, outputJsonObject.value);
    Assert.assertEquals("Hello", outputJsonObject.message);
}
 

@Test
public void testDew() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ModifyBytes());
    runner.setProperty(ModifyBytes.START_OFFSET, "94 B");
    runner.setProperty(ModifyBytes.END_OFFSET, "96 B");

    runner.enqueue(testFilePath);
    runner.run();

    runner.assertAllFlowFilesTransferred(ModifyBytes.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ModifyBytes.REL_SUCCESS).get(0);
    final String outContent = new String(out.toByteArray(), StandardCharsets.UTF_8);
    System.out.println(outContent);
    out.assertContentEquals("Dew".getBytes("UTF-8"));
}
 

@Test
public void test400WithPenalizeNoRetry() throws Exception {
    addHandler(new GetOrHeadHandler());

    runner.setProperty(InvokeHTTP.PROP_URL, url + "/status/400");
    runner.setProperty(InvokeHTTP.PROP_PENALIZE_NO_RETRY, "true");

    createFlowFiles(runner);

    runner.run();
    runner.assertTransferCount(InvokeHTTP.REL_SUCCESS_REQ, 0);
    runner.assertTransferCount(InvokeHTTP.REL_RESPONSE, 0);
    runner.assertTransferCount(InvokeHTTP.REL_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY, 1);
    runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
    runner.assertPenalizeCount(1);
    // expected in response
    final MockFlowFile bundle = runner.getFlowFilesForRelationship(InvokeHTTP.REL_NO_RETRY).get(0);
    final String actual = new String(bundle.toByteArray(), StandardCharsets.UTF_8);

    bundle.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "400");
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "Bad Request");
    bundle.assertAttributeEquals(InvokeHTTP.RESPONSE_BODY, "/status/400");
    final String expected = "Hello";
    Assert.assertEquals(expected, actual);
    bundle.assertAttributeEquals("Foo", "Bar");
}
 

@Test
public void testRidiculousRowCount() throws ClassNotFoundException, SQLException, InitializationException, IOException {
    long rowCount = Long.parseLong(Integer.toString(Integer.MAX_VALUE)) + 100;
    int partitionSize = 1000000;
    int expectedFileCount = (int) (rowCount / partitionSize) + 1;

    Connection conn = mock(Connection.class);
    when(dbcp.getConnection()).thenReturn(conn);
    Statement st = mock(Statement.class);
    when(conn.createStatement()).thenReturn(st);
    doNothing().when(st).close();
    ResultSet rs = mock(ResultSet.class);
    when(st.executeQuery(anyString())).thenReturn(rs);
    when(rs.next()).thenReturn(true);
    when(rs.getInt(1)).thenReturn((int) rowCount);
    when(rs.getLong(1)).thenReturn(rowCount);

    final ResultSetMetaData resultSetMetaData = mock(ResultSetMetaData.class);
    when(rs.getMetaData()).thenReturn(resultSetMetaData);
    when(resultSetMetaData.getColumnCount()).thenReturn(2);
    when(resultSetMetaData.getTableName(1)).thenReturn("");
    when(resultSetMetaData.getColumnType(1)).thenReturn(Types.INTEGER);
    when(resultSetMetaData.getColumnName(1)).thenReturn("COUNT");
    when(resultSetMetaData.getColumnType(2)).thenReturn(Types.INTEGER);
    when(resultSetMetaData.getColumnName(2)).thenReturn("ID");
    when(rs.getInt(2)).thenReturn(1000);


    runner.setProperty(GenerateTableFetch.TABLE_NAME, "TEST_QUERY_DB_TABLE");
    runner.setIncomingConnection(false);
    runner.setProperty(GenerateTableFetch.MAX_VALUE_COLUMN_NAMES, "ID");
    runner.setProperty(GenerateTableFetch.PARTITION_SIZE, Integer.toString(partitionSize));

    runner.run();
    runner.assertAllFlowFilesTransferred(REL_SUCCESS, expectedFileCount);
    MockFlowFile flowFile = runner.getFlowFilesForRelationship(REL_SUCCESS).get(0);
    String query = new String(flowFile.toByteArray());
    assertEquals("SELECT * FROM TEST_QUERY_DB_TABLE WHERE 1=1 ORDER BY ID FETCH NEXT 1000000 ROWS ONLY", query);
    runner.clearTransferState();
}
 

@Test
public void testRemoveHeader() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ModifyBytes());
    runner.setProperty(ModifyBytes.START_OFFSET, "12 B"); //REMOVE - '<<<HEADER>>>'
    runner.setProperty(ModifyBytes.END_OFFSET, "0 MB");

    runner.enqueue(testFilePath);
    runner.run();

    runner.assertAllFlowFilesTransferred(ModifyBytes.REL_SUCCESS, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(ModifyBytes.REL_SUCCESS).get(0);
    final String outContent = new String(out.toByteArray(), StandardCharsets.UTF_8);
    System.out.println(outContent);
    out.assertContentEquals(noHeaderFile);
}
 

private void assertOutputAvroRecords(List<Map<String, Object>> expectedRecords, MockFlowFile resultFlowFile) throws IOException {
    assertEquals(String.valueOf(expectedRecords.size()), resultFlowFile.getAttribute(PutHive3Streaming.HIVE_STREAMING_RECORD_COUNT_ATTR));

    final DataFileStream<GenericRecord> reader = new DataFileStream<>(
            new ByteArrayInputStream(resultFlowFile.toByteArray()),
            new GenericDatumReader<>());

    Schema schema = reader.getSchema();

    // Verify that the schema is preserved
    assertEquals(schema, new Schema.Parser().parse(new File("src/test/resources/user.avsc")));

    GenericRecord record = null;
    for (Map<String, Object> expectedRecord : expectedRecords) {
        assertTrue(reader.hasNext());
        record = reader.next(record);
        final String name = record.get("name").toString();
        final Integer favorite_number = (Integer) record.get("favorite_number");
        assertNotNull(name);
        assertNotNull(favorite_number);
        assertNull(record.get("favorite_color"));
        assertNull(record.get("scale"));

        assertEquals(expectedRecord.get("name"), name);
        assertEquals(expectedRecord.get("favorite_number"), favorite_number);
    }
    assertFalse(reader.hasNext());
}
 

@Test
public void testExecuteJar() throws Exception {
    File exJar = new File("src/test/resources/ExecuteCommand/TestSuccess.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.setValidateExpressionUsage(false);
    controller.enqueue(dummy.toPath());
    controller.setProperty(ExecuteStreamCommand.EXECUTION_COMMAND, "java");
    controller.setProperty(ExecuteStreamCommand.EXECUTION_ARGUMENTS, "-jar;" + jarPath);
    controller.run(1);
    controller.assertTransferCount(ExecuteStreamCommand.ORIGINAL_RELATIONSHIP, 1);
    controller.assertTransferCount(ExecuteStreamCommand.OUTPUT_STREAM_RELATIONSHIP, 1);

    List<MockFlowFile> flowFiles = controller.getFlowFilesForRelationship(ExecuteStreamCommand.OUTPUT_STREAM_RELATIONSHIP);
    MockFlowFile outputFlowFile = flowFiles.get(0);
    byte[] byteArray = outputFlowFile.toByteArray();
    String result = new String(byteArray);
    assertTrue(Pattern.compile("Test was a success\r?\n").matcher(result).find());
    assertEquals("0", outputFlowFile.getAttribute("execution.status"));
    assertEquals("java", outputFlowFile.getAttribute("execution.command"));
    assertEquals("-jar;", outputFlowFile.getAttribute("execution.command.args").substring(0, 5));
    String attribute = outputFlowFile.getAttribute("execution.command.args");
    String expected = "src" + File.separator + "test" + File.separator + "resources" + File.separator + "ExecuteCommand" + File.separator + "TestSuccess.jar";
    assertEquals(expected, attribute.substring(attribute.length() - expected.length()));

    MockFlowFile originalFlowFile = controller.getFlowFilesForRelationship(ExecuteStreamCommand.ORIGINAL_RELATIONSHIP).get(0);
    assertEquals(outputFlowFile.getAttribute("execution.status"), originalFlowFile.getAttribute("execution.status"));
    assertEquals(outputFlowFile.getAttribute("execution.command"), originalFlowFile.getAttribute("execution.command"));
    assertEquals(outputFlowFile.getAttribute("execution.command.args"), originalFlowFile.getAttribute("execution.command.args"));
}
 

public void parseFlowFiles(List<MockFlowFile> flowfiles) throws Exception, SAXException {
    for (MockFlowFile out : flowfiles) {
        final byte[] outData = out.toByteArray();
        final String outXml = new String(outData, "UTF-8");
        saxParser.parse(new InputSource(new StringReader(outXml)), new DefaultHandler());
    }
}
 

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

    runner.assertAllFlowFilesTransferred(ParseCEF.REL_SUCCESS, 1);
    final MockFlowFile mff = runner.getFlowFilesForRelationship(ParseCEF.REL_SUCCESS).get(0);

    byte [] rawJson = mff.toByteArray();

    JsonNode results = new ObjectMapper().readTree(rawJson);

    JsonNode header = results.get("header");
    JsonNode extension = results.get("extension");

    Assert.assertEquals("TestVendor", header.get("deviceVendor").asText());
    Assert.assertEquals(sdf.format(new Date(1423441663000L)),
                        extension.get("rt").asText());
    Assert.assertEquals("Test Long", extension.get("cn3Label").asText());
    Assert.assertEquals( 9223372036854775807L, extension.get("cn3").asLong());
    Assert.assertTrue(extension.get("cfp1").floatValue() == 1.234F);
    Assert.assertEquals("Test FP Number", extension.get("cfp1Label").asText());
    Assert.assertEquals("00:00:0c:07:ac:00", extension.get("smac").asText());
    Assert.assertEquals("2001:cdba:0:0:0:0:3257:9652", extension.get("c6a3").asText());
    Assert.assertEquals("Test IPv6", extension.get("c6a3Label").asText());
    Assert.assertEquals("123.123.123.123", extension.get("destinationTranslatedAddress").asText());
    Assert.assertEquals("Test String", extension.get("cs1Label").asText());
    Assert.assertEquals("test test test chocolate", extension.get("cs1").asText());
}
 

@Test
public void testSequenceFileBzipCompressionCodec() throws UnsupportedEncodingException, IOException {

    controller.setProperty(AbstractHadoopProcessor.COMPRESSION_CODEC, AbstractHadoopProcessor.CompressionType.BZIP.name());
    controller.setProperty(CreateHadoopSequenceFile.COMPRESSION_TYPE, SequenceFile.CompressionType.BLOCK.name());

    File inFile = inFiles[0];
    try (FileInputStream fin = new FileInputStream(inFile) ){
        controller.enqueue(fin);
    }
    controller.run();

    List<MockFlowFile> successSeqFiles = controller.getFlowFilesForRelationship(CreateHadoopSequenceFile.RELATIONSHIP_SUCCESS);
    List<MockFlowFile> failedFlowFiles = controller.getFlowFilesForRelationship(CreateHadoopSequenceFile.RELATIONSHIP_FAILURE);

    assertEquals(0, failedFlowFiles.size());
    assertEquals(1, successSeqFiles.size());

    MockFlowFile ff = successSeqFiles.iterator().next();
    byte[] data = ff.toByteArray();


    final String magicHeader = new String(data, 0, 3, "UTF-8");
    assertEquals("SEQ", magicHeader);
    // Format of header is SEQ followed by the version (1 byte).
    // Then, the length of the Key type (1 byte), then the Key type
    // Then, the length of the Value type(1 byte), then the Value type
    final String keyType = Text.class.getCanonicalName();
    final int valueTypeStart = 3 + 1 + 1 + keyType.length() + 1;
    final int valueTypeLength = data[5 + keyType.length()];
    final String valueType = BytesWritable.class.getCanonicalName();

    assertEquals(valueType.length(), valueTypeLength);
    assertEquals(valueType, new String(data, valueTypeStart, valueType.length(), "UTF-8"));

    final int compressionIndex = 3 + 1 + 1 + keyType.length() + 1 + valueType.length();
    final int blockCompressionIndex = compressionIndex + 1;

    assertEquals(1, data[compressionIndex]);
    assertEquals(1, data[blockCompressionIndex]);

    final int codecTypeSize = data[blockCompressionIndex + 1];
    final int codecTypeStartIndex = blockCompressionIndex + 2;

    assertEquals(BZip2Codec.class.getCanonicalName(), new String(data, codecTypeStartIndex, codecTypeSize, "UTF-8"));
}
 

@Test
public void test200DigestAuth() throws Exception {
    addHandler(new DigestAuthHandler());
    final String username = "basic_user";
    final String password = "basic_password";

    runner.setProperty(InvokeHTTP.PROP_URL, url + "/status/200");
    runner.setProperty(InvokeHTTP.PROP_BASIC_AUTH_USERNAME, username);
    runner.setProperty(InvokeHTTP.PROP_BASIC_AUTH_PASSWORD, password);
    runner.setProperty(InvokeHTTP.PROP_DIGEST_AUTH,"true");

    createFlowFiles(runner);

    runner.run();

    runner.assertTransferCount(InvokeHTTP.REL_SUCCESS_REQ, 1);
    runner.assertTransferCount(InvokeHTTP.REL_RESPONSE, 1);
    runner.assertTransferCount(InvokeHTTP.REL_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_NO_RETRY, 0);
    runner.assertTransferCount(InvokeHTTP.REL_FAILURE, 0);
    runner.assertPenalizeCount(0);

    //expected in request status.code and status.message
    //original flow file (+attributes)
    final MockFlowFile bundle = runner.getFlowFilesForRelationship(InvokeHTTP.REL_SUCCESS_REQ).get(0);
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
    bundle.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
    bundle.assertAttributeEquals("Foo", "Bar");
    final String actual = new String(bundle.toByteArray(), StandardCharsets.UTF_8);
    final String expected = "Hello";
    Assert.assertEquals(expected, actual);

    //expected in response
    //status code, status message, all headers from server response --> ff attributes
    //server response message body into payload of ff
    final MockFlowFile bundle1 = runner.getFlowFilesForRelationship(InvokeHTTP.REL_RESPONSE).get(0);
    bundle1.assertContentEquals(("DIGEST"+System.lineSeparator()).getBytes("UTF-8"));
    bundle1.assertAttributeEquals(InvokeHTTP.STATUS_CODE, "200");
    bundle1.assertAttributeEquals(InvokeHTTP.STATUS_MESSAGE, "OK");
    bundle1.assertAttributeEquals("Foo", "Bar");
    bundle1.assertAttributeEquals("Content-Type", "text/plain;charset=iso-8859-1");
}
 

public void invokeOnTrigger(final String query, final boolean incomingFlowFile, String outputFormat,
        String preQueries, String postQueries)
        throws InitializationException, ClassNotFoundException, SQLException, IOException {

    TestRunner runner = doOnTrigger(query, incomingFlowFile, outputFormat, preQueries, postQueries);
    runner.assertAllFlowFilesTransferred(SelectHiveQL.REL_SUCCESS, 1);

    final List<MockFlowFile> flowfiles = runner.getFlowFilesForRelationship(SelectHiveQL.REL_SUCCESS);
    MockFlowFile flowFile = flowfiles.get(0);
    final InputStream in = new ByteArrayInputStream(flowFile.toByteArray());
    long recordsFromStream = 0;
    if (AVRO.equals(outputFormat)) {
        assertEquals(MIME_TYPE_AVRO_BINARY, flowFile.getAttribute(CoreAttributes.MIME_TYPE.key()));
        final DatumReader<GenericRecord> datumReader = new GenericDatumReader<>();
        try (DataFileStream<GenericRecord> dataFileReader = new DataFileStream<>(in, datumReader)) {
            GenericRecord record = null;
            while (dataFileReader.hasNext()) {
                // Reuse record object by passing it to next(). This saves us from
                // allocating and garbage collecting many objects for files with
                // many items.
                record = dataFileReader.next(record);
                recordsFromStream++;
            }
        }
    } else {
        assertEquals(CSV_MIME_TYPE, flowFile.getAttribute(CoreAttributes.MIME_TYPE.key()));
        BufferedReader br = new BufferedReader(new InputStreamReader(in));

        String headerRow = br.readLine();
        // Derby capitalizes column names
        assertEquals("PERSONID,PERSONNAME,PERSONCODE", headerRow);

        // Validate rows
        String line;
        while ((line = br.readLine()) != null) {
            recordsFromStream++;
            String[] values = line.split(",");
            if (recordsFromStream < (NUM_OF_ROWS - 10)) {
                assertEquals(3, values.length);
                assertTrue(values[1].startsWith("\""));
                assertTrue(values[1].endsWith("\""));
            } else {
                assertEquals(2, values.length); // Middle value is null
            }
        }
    }
    assertEquals(NUM_OF_ROWS - 10, recordsFromStream);
    assertEquals(recordsFromStream, Integer.parseInt(flowFile.getAttribute(SelectHiveQL.RESULT_ROW_COUNT)));
    flowFile.assertAttributeEquals(AbstractHiveQLProcessor.ATTR_INPUT_TABLES, "persons");
}
 

@Test
public void testColumnTypeMissing() throws ClassNotFoundException, SQLException, InitializationException, IOException {
    // Load test data to database
    final Connection con = ((DBCPService) runner.getControllerService("dbcp")).getConnection();
    Statement stmt = con.createStatement();

    try {
        stmt.execute("drop table TEST_QUERY_DB_TABLE");
    } catch (final SQLException sqle) {
        // Ignore this error, probably a "table does not exist" since Derby doesn't yet support DROP IF EXISTS [DERBY-4842]
    }

    stmt.execute("create table TEST_QUERY_DB_TABLE (id integer not null, bucket integer not null)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (0, 0)");
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (1, 0)");

    runner.setProperty(GenerateTableFetch.TABLE_NAME, "TEST_QUERY_DB_TABLE");
    runner.setIncomingConnection(true);

    runner.setProperty(GenerateTableFetch.TABLE_NAME, "${tableName}");
    runner.setIncomingConnection(true);
    runner.setProperty(GenerateTableFetch.MAX_VALUE_COLUMN_NAMES, "${maxValueCol}");
    runner.enqueue("".getBytes(), new HashMap<String, String>() {{
        put("tableName", "TEST_QUERY_DB_TABLE");
        put("maxValueCol", "id");
    }});
    runner.run();
    runner.assertAllFlowFilesTransferred(REL_SUCCESS, 1);
    MockFlowFile flowFile = runner.getFlowFilesForRelationship(REL_SUCCESS).get(0);
    String query = new String(flowFile.toByteArray());
    assertEquals("SELECT * FROM TEST_QUERY_DB_TABLE WHERE id <= 1 ORDER BY id FETCH NEXT 10000 ROWS ONLY", query);
    runner.clearTransferState();


    // Clear columnTypeMap to simulate it's clean after instance reboot
    processor.columnTypeMap.clear();

    // Insert new records
    stmt.execute("insert into TEST_QUERY_DB_TABLE (id, bucket) VALUES (2, 0)");

    // Re-launch FlowFile to se if re-cache column type works
    runner.enqueue("".getBytes(), new HashMap<String, String>() {{
        put("tableName", "TEST_QUERY_DB_TABLE");
        put("maxValueCol", "id");
    }});

    // It should re-cache column type
    runner.run();
    runner.assertAllFlowFilesTransferred(REL_SUCCESS, 1);
    flowFile = runner.getFlowFilesForRelationship(REL_SUCCESS).get(0);
    query = new String(flowFile.toByteArray());
    assertEquals("SELECT * FROM TEST_QUERY_DB_TABLE WHERE id > 1 AND id <= 2 ORDER BY id FETCH NEXT 10000 ROWS ONLY", query);
    runner.clearTransferState();
}
 

/**
 * If the output fields are added to a non-record field, then the result should be that the field
 * becomes a UNION that does allow the Record and the value is set to a Record.
 */
@Test
public void testAddFieldsToNonRecordField() throws InitializationException {
    final RecordLookup lookupService = new RecordLookup();
    runner.addControllerService("lookup", lookupService);
    runner.enableControllerService(lookupService);

    final List<RecordField> fields = new ArrayList<>();
    fields.add(new RecordField("favorite", RecordFieldType.STRING.getDataType()));
    fields.add(new RecordField("least", RecordFieldType.STRING.getDataType()));
    final RecordSchema schema = new SimpleRecordSchema(fields);
    final Record sports = new MapRecord(schema, new HashMap<>());

    sports.setValue("favorite", "basketball");
    sports.setValue("least", "soccer");

    lookupService.addValue("John Doe", sports);

    recordReader = new MockRecordParser();
    recordReader.addSchemaField("name", RecordFieldType.STRING);
    recordReader.addSchemaField("age", RecordFieldType.INT);
    recordReader.addSchemaField("sport", RecordFieldType.STRING);

    recordReader.addRecord("John Doe", 48, null);

    runner.addControllerService("reader", recordReader);
    runner.enableControllerService(recordReader);

    runner.setProperty("lookup", "/name");
    runner.setProperty(LookupRecord.RESULT_RECORD_PATH, "/sport");
    runner.setProperty(LookupRecord.RESULT_CONTENTS, LookupRecord.RESULT_RECORD_FIELDS);

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

    final MockFlowFile out = runner.getFlowFilesForRelationship(LookupRecord.REL_MATCHED).get(0);

    // We can't be sure of the order of the fields in the record, so we allow either 'least' or 'favorite' to be first
    final String outputContents = new String(out.toByteArray());
    assertTrue(outputContents.equals("John Doe,48,MapRecord[{favorite=basketball, least=soccer}]\n")
        || outputContents.equals("John Doe,48,MapRecord[{least=soccer, favorite=basketball}]\n"));
}