Examples with MockFlowFile org.apache.nifi.util.MockFlowFile used on opensource projects

Example 66 with MockFlowFile

use of org.apache.nifi.util.MockFlowFile in project nifi by apache.

the class TestExtractAvroMetadata method testExtractionWithMD5.

@Test
public void testExtractionWithMD5() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractAvroMetadata());
    runner.setProperty(ExtractAvroMetadata.FINGERPRINT_ALGORITHM, ExtractAvroMetadata.MD5);
    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/user.avsc"));
    final ByteArrayOutputStream out = getOutputStreamWithOneUser(schema);
    runner.enqueue(out.toByteArray());
    runner.run();
    runner.assertAllFlowFilesTransferred(ExtractAvroMetadata.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(ExtractAvroMetadata.REL_SUCCESS).get(0);
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_FINGERPRINT_ATTR, "3c6a7bee8994be20314dd28c6a3af4f2");
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_TYPE_ATTR, Schema.Type.RECORD.getName());
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_NAME_ATTR, "User");
    flowFile.assertAttributeNotExists(AVRO_SCHEMA_ATTR);
}

Example 67 with MockFlowFile

use of org.apache.nifi.util.MockFlowFile in project nifi by apache.

the class TestExtractAvroMetadata method testExtractionWithNonRecordSchema.

@Test
public void testExtractionWithNonRecordSchema() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractAvroMetadata());
    runner.setProperty(ExtractAvroMetadata.COUNT_ITEMS, "true");
    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/array.avsc"));
    final GenericData.Array<String> data = new GenericData.Array<>(schema, Arrays.asList("one", "two", "three"));
    final DatumWriter<GenericData.Array<String>> datumWriter = new GenericDatumWriter<>(schema);
    final ByteArrayOutputStream out = new ByteArrayOutputStream();
    final DataFileWriter<GenericData.Array<String>> dataFileWriter = new DataFileWriter<>(datumWriter);
    dataFileWriter.create(schema, out);
    dataFileWriter.append(data);
    dataFileWriter.append(data);
    dataFileWriter.close();
    runner.enqueue(out.toByteArray());
    runner.run();
    runner.assertAllFlowFilesTransferred(ExtractAvroMetadata.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(ExtractAvroMetadata.REL_SUCCESS).get(0);
    flowFile.assertAttributeExists(ExtractAvroMetadata.SCHEMA_FINGERPRINT_ATTR);
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_TYPE_ATTR, Schema.Type.ARRAY.getName());
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_NAME_ATTR, "array");
    // number of arrays, not elements
    flowFile.assertAttributeEquals(ExtractAvroMetadata.ITEM_COUNT_ATTR, "2");
}

Example 68 with MockFlowFile

use of org.apache.nifi.util.MockFlowFile in project nifi by apache.

the class TestExtractAvroMetadata method testDefaultExtraction.

@Test
public void testDefaultExtraction() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new ExtractAvroMetadata());
    final Schema schema = new Schema.Parser().parse(new File("src/test/resources/user.avsc"));
    final ByteArrayOutputStream out = getOutputStreamWithOneUser(schema);
    runner.enqueue(out.toByteArray());
    runner.run();
    runner.assertAllFlowFilesTransferred(ExtractAvroMetadata.REL_SUCCESS, 1);
    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(ExtractAvroMetadata.REL_SUCCESS).get(0);
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_FINGERPRINT_ATTR, "b2d1d8d3de2833ce");
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_TYPE_ATTR, Schema.Type.RECORD.getName());
    flowFile.assertAttributeEquals(ExtractAvroMetadata.SCHEMA_NAME_ATTR, "User");
    flowFile.assertAttributeNotExists(AVRO_SCHEMA_ATTR);
    flowFile.assertAttributeNotExists(ExtractAvroMetadata.ITEM_COUNT_ATTR);
}

Example 69 with MockFlowFile

use of org.apache.nifi.util.MockFlowFile in project nifi by apache.

the class TestSplitAvro method testRecordSplitDatafileOutputWithoutMetadata.

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

Example 70 with MockFlowFile

use of org.apache.nifi.util.MockFlowFile in project nifi by apache.

the class TestSplitAvro method checkBareRecordsSplitSize.

private void checkBareRecordsSplitSize(final List<MockFlowFile> flowFiles, final int expectedRecordsPerSplit, final boolean checkMetadata) throws IOException {
    for (final MockFlowFile flowFile : flowFiles) {
        try (final ByteArrayInputStream in = new ByteArrayInputStream(flowFile.toByteArray())) {
            final DatumReader<GenericRecord> reader = new GenericDatumReader<>(schema);
            final Decoder decoder = DecoderFactory.get().binaryDecoder(in, null);
            int count = 0;
            GenericRecord record = reader.read(null, decoder);
            try {
                while (record != null) {
                    Assert.assertNotNull(record.get("name"));
                    Assert.assertNotNull(record.get("favorite_number"));
                    count++;
                    record = reader.read(record, decoder);
                }
            } catch (EOFException eof) {
            // expected
            }
            assertEquals(expectedRecordsPerSplit, count);
        }
        if (checkMetadata) {
            Assert.assertTrue(flowFile.getAttributes().containsKey(META_KEY1));
            Assert.assertTrue(flowFile.getAttributes().containsKey(META_KEY2));
            Assert.assertTrue(flowFile.getAttributes().containsKey(META_KEY3));
        }
    }
}