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

Example 71 with MockFlowFile

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

the class TestSplitAvro method testRecordSplitDatafileOutputWithSplitSizeLarger.

@Test
public void testRecordSplitDatafileOutputWithSplitSizeLarger() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new SplitAvro());
    runner.setProperty(SplitAvro.OUTPUT_SIZE, "200");
    runner.enqueue(users.toByteArray());
    runner.run();
    runner.assertTransferCount(SplitAvro.REL_SPLIT, 1);
    runner.assertTransferCount(SplitAvro.REL_ORIGINAL, 1);
    runner.assertTransferCount(SplitAvro.REL_FAILURE, 0);
    runner.getFlowFilesForRelationship(SplitAvro.REL_ORIGINAL).get(0).assertAttributeEquals(FRAGMENT_COUNT.key(), "1");
    final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(SplitAvro.REL_SPLIT);
    checkDataFileSplitSize(flowFiles, 100, true);
}

Example 72 with MockFlowFile

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

the class TestSplitAvro method checkDataFileTotalSize.

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

Example 73 with MockFlowFile

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

the class ITDeleteElasticsearch5Test method testDeleteIntegrationTestDocumentNotFound.

@Test
public void testDeleteIntegrationTestDocumentNotFound() {
    final TestRunner runnerDelete = TestRunners.newTestRunner(new DeleteElasticsearch5());
    runnerDelete.setValidateExpressionUsage(false);
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.CLUSTER_NAME, clusterName);
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.HOSTS, "127.0.0.1:9300");
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.PING_TIMEOUT, "5s");
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.SAMPLER_INTERVAL, "5s");
    runnerDelete.setProperty(DeleteElasticsearch5.INDEX, INDEX1);
    runnerDelete.setProperty(DeleteElasticsearch5.TYPE, TYPE1);
    runnerDelete.setProperty(DeleteElasticsearch5.DOCUMENT_ID, "${documentId}");
    runnerDelete.assertValid();
    runnerDelete.enqueue(new byte[] {}, new HashMap<String, String>() {

        {
            put("documentId", documentId);
        }
    });
    runnerDelete.enqueue(new byte[] {});
    runnerDelete.run(1, true, true);
    runnerDelete.assertAllFlowFilesTransferred(DeleteElasticsearch5.REL_NOT_FOUND, 1);
    final MockFlowFile out = runnerDelete.getFlowFilesForRelationship(DeleteElasticsearch5.REL_NOT_FOUND).get(0);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_FILENAME, documentId);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_INDEX, INDEX1);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_TYPE, TYPE1);
}

Example 74 with MockFlowFile

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

the class ITDeleteElasticsearch5Test method testDeleteIntegrationTestBadType.

@Test
public void testDeleteIntegrationTestBadType() {
    final TestRunner runnerDelete = TestRunners.newTestRunner(new DeleteElasticsearch5());
    runnerDelete.setValidateExpressionUsage(false);
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.CLUSTER_NAME, clusterName);
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.HOSTS, "127.0.0.1:9300");
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.PING_TIMEOUT, "5s");
    runnerDelete.setProperty(AbstractElasticsearch5TransportClientProcessor.SAMPLER_INTERVAL, "5s");
    runnerDelete.setProperty(DeleteElasticsearch5.INDEX, INDEX1);
    String type = String.valueOf(System.currentTimeMillis());
    runnerDelete.setProperty(DeleteElasticsearch5.TYPE, type);
    runnerDelete.setProperty(DeleteElasticsearch5.DOCUMENT_ID, "${documentId}");
    runnerDelete.assertValid();
    runnerDelete.enqueue(new byte[] {}, new HashMap<String, String>() {

        {
            put("documentId", documentId);
        }
    });
    runnerDelete.enqueue(new byte[] {});
    runnerDelete.run(1, true, true);
    runnerDelete.assertAllFlowFilesTransferred(DeleteElasticsearch5.REL_NOT_FOUND, 1);
    final MockFlowFile out = runnerDelete.getFlowFilesForRelationship(DeleteElasticsearch5.REL_NOT_FOUND).get(0);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_FILENAME, documentId);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_INDEX, INDEX1);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_TYPE, type);
}

Example 75 with MockFlowFile

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

the class TestDeleteElasticsearch5 method testDeleteRetryableException.

@Test
public void testDeleteRetryableException() throws IOException {
    mockDeleteProcessor = new DeleteElasticsearch5() {

        @Override
        protected DeleteRequestBuilder prepareDeleteRequest(String index, String docId, String docType) {
            return null;
        }

        @Override
        protected DeleteResponse doDelete(DeleteRequestBuilder requestBuilder) throws InterruptedException, ExecutionException {
            throw new ElasticsearchTimeoutException("timeout");
        }

        @Override
        public void setup(ProcessContext context) {
        }
    };
    runner = TestRunners.newTestRunner(mockDeleteProcessor);
    runner.setValidateExpressionUsage(true);
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.CLUSTER_NAME, "elasticsearch");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.HOSTS, "127.0.0.1:9300");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.PING_TIMEOUT, "5s");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.SAMPLER_INTERVAL, "5s");
    runner.setProperty(DeleteElasticsearch5.INDEX, INDEX1);
    runner.setProperty(DeleteElasticsearch5.TYPE, TYPE1);
    runner.setProperty(DeleteElasticsearch5.DOCUMENT_ID, "${documentId}");
    runner.assertValid();
    runner.enqueue(new byte[] {}, new HashMap<String, String>() {

        {
            put("documentId", documentId);
        }
    });
    runner.run(1, true, true);
    runner.assertAllFlowFilesTransferred(DeleteElasticsearch5.REL_RETRY, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(DeleteElasticsearch5.REL_RETRY).get(0);
    assertNotNull(out);
    assertEquals("timeout", out.getAttribute(DeleteElasticsearch5.ES_ERROR_MESSAGE));
    out.assertAttributeEquals(DeleteElasticsearch5.ES_REST_STATUS, null);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_FILENAME, documentId);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_INDEX, INDEX1);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_TYPE, TYPE1);
}