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

@Test(expected = AssertionError.class)
public void testPutElasticSearchBadHostInEL() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new PutElasticsearchHttpRecord());

    runner.setProperty(AbstractElasticsearchHttpProcessor.ES_URL, "${es.url}");
    runner.setProperty(PutElasticsearchHttpRecord.INDEX, "doc");
    runner.setProperty(PutElasticsearchHttpRecord.TYPE, "status");
    runner.setProperty(PutElasticsearchHttpRecord.ID_RECORD_PATH, "/id");
    runner.assertValid();

    runner.enqueue(new byte[0], new HashMap<String, String>() {{
        put("doc_id", "1");
    }});

    runner.run();
}
 

@Test
public void testSimpleDeleteWithFilename() throws Exception {
    putTestFile(KEY, new byte[]{7, 8, 9});
    assertTrue(fileExists(KEY));

    final TestRunner runner = buildNewRunner(new DeleteGCSObject());
    runner.setProperty(DeleteGCSObject.BUCKET, BUCKET);
    runner.assertValid();

    runner.enqueue("testdata", ImmutableMap.of(
            CoreAttributes.FILENAME.key(), KEY
    ));

    runner.run();

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

    assertFalse(fileExists(KEY));
}
 

@Test(expected = IllegalArgumentException.class)
public void testScanWithInvalidFilter() throws InitializationException, IOException {
    final String tableName = "nifi";
    final TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);

    // Mock an HBase Table so we can verify the put operations later
    final Table table = Mockito.mock(Table.class);
    when(table.getName()).thenReturn(TableName.valueOf(tableName));

    // create the controller service and link it to the test processor
    final MockHBaseClientService service = configureHBaseClientService(runner, table);
    runner.assertValid(service);

    // perform a scan and verify the four rows were returned
    final CollectingResultHandler handler = new CollectingResultHandler();
    final HBaseClientService hBaseClientService = runner.getProcessContext().getProperty(TestProcessor.HBASE_CLIENT_SERVICE)
            .asControllerService(HBaseClientService.class);

    // this should throw IllegalArgumentException
    final String filter = "this is not a filter";
    hBaseClientService.scan(tableName, new ArrayList<Column>(), filter, System.currentTimeMillis(), handler);
}
 

@Test
public void testJsonStringCredentials() throws Exception {
    final String jsonRead = new String(
            Files.readAllBytes(Paths.get("src/test/resources/mock-gcp-service-account.json"))
    );
    final TestRunner runner = TestRunners.newTestRunner(MockCredentialsFactoryProcessor.class);
    runner.setProperty(CredentialPropertyDescriptors.SERVICE_ACCOUNT_JSON,
            jsonRead);
    runner.assertValid();

    Map<PropertyDescriptor, String> properties = runner.getProcessContext().getProperties();
    final CredentialsFactory factory = new CredentialsFactory();
    final GoogleCredentials credentials = factory.getGoogleCredentials(properties);

    assertNotNull(credentials);
    assertEquals("credentials class should be equal", ServiceAccountCredentials.class,
            credentials.getClass());
}
 

@Test
public void testTimestampExpressionInvalidRoutesToFailure() throws Exception {
    MockPutCloudWatchMetric mockPutCloudWatchMetric = new MockPutCloudWatchMetric();
    final TestRunner runner = TestRunners.newTestRunner(mockPutCloudWatchMetric);

    runner.setProperty(PutCloudWatchMetric.NAMESPACE, "TestNamespace");
    runner.setProperty(PutCloudWatchMetric.METRIC_NAME, "TestMetric");
    runner.setProperty(PutCloudWatchMetric.UNIT, "Count");
    runner.setProperty(PutCloudWatchMetric.VALUE, "1");
    runner.setProperty(PutCloudWatchMetric.TIMESTAMP, "${timestamp.value}");
    runner.assertValid();

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("timestamp.value", "1476296132575broken");
    runner.enqueue(new byte[] {}, attributes);
    runner.run();

    Assert.assertEquals(0, mockPutCloudWatchMetric.putMetricDataCallCount);
    runner.assertAllFlowFilesTransferred(PutCloudWatchMetric.REL_FAILURE, 1);
}
 

@Test
public void testHttpsUrlShouldRequireSSLContext() throws InitializationException {
    final TestRunner runner = TestRunners.newTestRunner(PutSolrRecord.class);
    MockRecordParser recordParser = new MockRecordParser();
    recordParser.addRecord(1, "Abhinav","R",8,"Chemistry","term1", 98);
    runner.addControllerService("parser", recordParser);
    runner.enableControllerService(recordParser);
    runner.setProperty(PutSolrRecord.RECORD_READER, "parser");

    runner.setProperty(SolrUtils.SOLR_TYPE, SolrUtils.SOLR_TYPE_STANDARD.getValue());
    runner.setProperty(SolrUtils.SOLR_LOCATION, "https://localhost:8443/solr");
    runner.assertNotValid();

    final SSLContextService sslContextService = new MockSSLContextService();
    runner.addControllerService("ssl-context", sslContextService);
    runner.enableControllerService(sslContextService);

    runner.setProperty(SolrUtils.SSL_CONTEXT_SERVICE, "ssl-context");
    runner.assertValid();
}
 

@Test
public void testIllegalOrderValue() {
    final TestRunner runner = TestRunners.newTestRunner(EnforceOrder.class);

    runner.setProperty(EnforceOrder.GROUP_IDENTIFIER, "${group}");
    runner.setProperty(EnforceOrder.ORDER_ATTRIBUTE, "index");
    runner.setProperty(EnforceOrder.INITIAL_ORDER, "1");
    runner.assertValid();
    Ordered.enqueue(runner, "b", 1);
    Ordered.enqueue(runner, "a", 2);
    runner.enqueue("illegal order", Ordered.i("a", 1).put("index", "non-integer").map());
    Ordered.enqueue(runner, "a", 1);

    runner.run();

    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(EnforceOrder.REL_SUCCESS);
    assertEquals(3, succeeded.size());
    succeeded.sort(new FirstInFirstOutPrioritizer());
    succeeded.get(0).assertContentEquals("a.1");
    succeeded.get(1).assertContentEquals("a.2");
    succeeded.get(2).assertContentEquals("b.1");

    final List<MockFlowFile> failed = runner.getFlowFilesForRelationship(EnforceOrder.REL_FAILURE);
    assertEquals(1, failed.size());
    failed.get(0).assertAttributeExists(EnforceOrder.ATTR_DETAIL);
    failed.get(0).assertContentEquals("illegal order");
}
 

@Test
public void testInvalidRegexInAttribute() {
    final TestRunner runner = TestRunners.newTestRunner(new UpdateAttribute());
    runner.setProperty(UpdateAttribute.DELETE_ATTRIBUTES, "${butter}");
    runner.assertValid();

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("butter", "(");
    runner.enqueue(new byte[0], attributes);
    try {
        runner.run();
    } catch (Throwable t) {
        assertEquals(t.getCause().getClass(), PatternSyntaxException.class);
    }
}
 

@Test
@Ignore("Comment this out if you want to run against local or test ES")
public void testFetchElasticsearch5Batch() throws IOException {
    System.out.println("Starting test " + new Object() {
    }.getClass().getEnclosingMethod().getName());
    final TestRunner runner = TestRunners.newTestRunner(new FetchElasticsearch5());
    runner.setValidateExpressionUsage(true);

    //Local Cluster - Mac pulled from brew
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.CLUSTER_NAME, "elasticsearch_brew");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.HOSTS, "127.0.0.1:9300");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.PING_TIMEOUT, "5s");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.SAMPLER_INTERVAL, "5s");
    runner.setProperty(FetchElasticsearch5.INDEX, "doc");

    runner.setProperty(FetchElasticsearch5.TYPE, "status");
    runner.setProperty(FetchElasticsearch5.DOC_ID, "${doc_id}");
    runner.assertValid();


    String message = convertStreamToString(docExample);
    for (int i = 0; i < 100; i++) {

        long newId = 28039652140L + i;
        final String newStrId = Long.toString(newId);
        runner.enqueue(message.getBytes(), new HashMap<String, String>() {{
            put("doc_id", newStrId);
        }});

    }

    runner.run();

    runner.assertAllFlowFilesTransferred(FetchElasticsearch5.REL_SUCCESS, 100);
}
 

@Test
public void testGlobDelete() throws Exception {
    Path glob = new Path("/data/for/2017/08/05/*");
    int fileCount = 300;
    FileStatus[] fileStatuses = new FileStatus[fileCount];
    for (int i = 0; i < fileCount; i++) {
        Path file = new Path("/data/for/2017/08/05/file" + i);
        FileStatus fileStatus = mock(FileStatus.class);
        when(fileStatus.getPath()).thenReturn(file);
        fileStatuses[i] = fileStatus;
    }
    when(mockFileSystem.exists(any(Path.class))).thenReturn(true);
    when(mockFileSystem.globStatus(any(Path.class))).thenReturn(fileStatuses);
    DeleteHDFS deleteHDFS = new TestableDeleteHDFS(kerberosProperties, mockFileSystem);
    TestRunner runner = TestRunners.newTestRunner(deleteHDFS);
    runner.setIncomingConnection(false);
    runner.assertNotValid();
    runner.setProperty(DeleteHDFS.FILE_OR_DIRECTORY, glob.toString());
    runner.assertValid();
    runner.run();
    runner.assertAllFlowFilesTransferred(DeleteHDFS.REL_SUCCESS);
    runner.assertTransferCount(DeleteHDFS.REL_SUCCESS, fileCount);
    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(DeleteHDFS.REL_SUCCESS);
    for (int i = 0; i < fileCount; i++) {
        FlowFile flowFile = flowFiles.get(i);
        assertEquals("file" + i, flowFile.getAttribute("filename"));
        assertEquals("/data/for/2017/08/05", flowFile.getAttribute("path"));
    }
}
 

@Test
public void testAssumeRoleCredentials() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
    runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
    runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_ARN, "BogusArn");
    runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_NAME, "BogusSession");
    runner.assertValid();

    Map<PropertyDescriptor, String> properties = runner.getProcessContext().getProperties();
    final CredentialsProviderFactory factory = new CredentialsProviderFactory();
    final AWSCredentialsProvider credentialsProvider = factory.getCredentialsProvider(properties);
    Assert.assertNotNull(credentialsProvider);
    assertEquals("credentials provider should be equal", STSAssumeRoleSessionCredentialsProvider.class,
            credentialsProvider.getClass());
}
 

@Test
public void testBasicState() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(new UpdateAttribute());
    runner.setProperty(UpdateAttribute.STORE_STATE, STORE_STATE_LOCALLY);
    runner.setProperty("count", "${getStateValue('count'):plus(1)}");
    runner.setProperty("sum", "${getStateValue('sum'):plus(${pencils})}");

    runner.assertNotValid();
    runner.setProperty(UpdateAttribute.STATEFUL_VARIABLES_INIT_VALUE, "0");
    runner.assertValid();

    final Map<String, String> attributes2 = new HashMap<>();
    attributes2.put("pencils", "2");

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

    final Map<String, String> attributes3 = new HashMap<>();
    attributes3.put("pencils", "3");
    runner.enqueue(new byte[0], attributes3);
    runner.enqueue(new byte[0], attributes3);

    final Map<String, String> attributes5 = new HashMap<>();
    attributes5.put("pencils", "5");
    runner.enqueue(new byte[0], attributes5);

    runner.run(5);

    runner.assertAllFlowFilesTransferred(UpdateAttribute.REL_SUCCESS, 5);
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_SUCCESS).get(4).assertAttributeEquals("count", "5");
    runner.getFlowFilesForRelationship(UpdateAttribute.REL_SUCCESS).get(4).assertAttributeEquals("sum", "15");
}
 

@Test
public void testAssumeRoleCredentials() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
    runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
    runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_ARN, "BogusArn");
    runner.setProperty(CredentialPropertyDescriptors.ASSUME_ROLE_NAME, "BogusSession");
    runner.assertValid();

    Map<PropertyDescriptor, String> properties = runner.getProcessContext().getProperties();
    final CredentialsProviderFactory factory = new CredentialsProviderFactory();
    final AWSCredentialsProvider credentialsProvider = factory.getCredentialsProvider(properties);
    Assert.assertNotNull(credentialsProvider);
    assertEquals("credentials provider should be equal", STSAssumeRoleSessionCredentialsProvider.class,
            credentialsProvider.getClass());
}
 

protected TestRunner setCredentialsCS(TestRunner runner) throws InitializationException {
    final String serviceAccountJsonFilePath = "path/to/credentials/json";
    final Map<String, String> propertiesMap = new HashMap<>();
    final GCPCredentialsControllerService credentialsControllerService = new GCPCredentialsControllerService();

    propertiesMap.put("application-default-credentials", "false");
    propertiesMap.put("compute-engine-credentials", "false");
    propertiesMap.put("service-account-json-file", serviceAccountJsonFilePath);

    runner.addControllerService(CONTROLLER_SERVICE, credentialsControllerService, propertiesMap);
    runner.enableControllerService(credentialsControllerService);
    runner.assertValid(credentialsControllerService);

    return runner;
}
 

@Test
public void testMoveAndKeep() throws IOException {
    final File sourceFile = new File("target/1.txt");
    final byte[] content = "Hello, World!".getBytes();
    Files.write(sourceFile.toPath(), content, StandardOpenOption.CREATE);

    final TestRunner runner = TestRunners.newTestRunner(new FetchFile());
    runner.setProperty(FetchFile.FILENAME, sourceFile.getAbsolutePath());
    runner.setProperty(FetchFile.COMPLETION_STRATEGY, FetchFile.COMPLETION_MOVE.getValue());
    runner.assertNotValid();
    runner.setProperty(FetchFile.MOVE_DESTINATION_DIR, "target/move-target");
    runner.setProperty(FetchFile.CONFLICT_STRATEGY, FetchFile.CONFLICT_KEEP_INTACT.getValue());
    runner.assertValid();

    final File destDir = new File("target/move-target");
    final File destFile = new File(destDir, sourceFile.getName());

    final byte[] goodBye = "Good-bye".getBytes();
    Files.write(destFile.toPath(), goodBye);

    runner.enqueue(new byte[0]);
    runner.run();
    runner.assertAllFlowFilesTransferred(FetchFile.REL_SUCCESS, 1);
    runner.getFlowFilesForRelationship(FetchFile.REL_SUCCESS).get(0).assertContentEquals(content);

    final byte[] replacedContent = Files.readAllBytes(destFile.toPath());
    assertTrue(Arrays.equals(goodBye, replacedContent));
    assertFalse(sourceFile.exists());
    assertTrue(destFile.exists());
}
 

@Test
public void testFileCredentials() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(MockAWSProcessor.class);
    runner.setProperty(CredentialPropertyDescriptors.CREDENTIALS_FILE, "src/test/resources/mock-aws-credentials.properties");
    runner.assertValid();

    Map<PropertyDescriptor, String> properties = runner.getProcessContext().getProperties();
    final CredentialsProviderFactory factory = new CredentialsProviderFactory();
    final AWSCredentialsProvider credentialsProvider = factory.getCredentialsProvider(properties);
    Assert.assertNotNull(credentialsProvider);
    assertEquals("credentials provider should be equal", PropertiesFileCredentialsProvider.class,
            credentialsProvider.getClass());
}
 

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

    runner.removeProperty(FetchGCSObject.KEY);
    runner.removeProperty(FetchGCSObject.BUCKET);

    runner.setProperty(FetchGCSObject.GENERATION, String.valueOf(GENERATION));
    runner.assertValid();

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

    runner.enqueue("", ImmutableMap.of(
            BUCKET_ATTR, BUCKET,
            CoreAttributes.FILENAME.key(), KEY
    ));

    runner.run();

    ArgumentCaptor<BlobId> blobIdArgumentCaptor = ArgumentCaptor.forClass(BlobId.class);
    ArgumentCaptor<Storage.BlobSourceOption> blobSourceOptionArgumentCaptor = ArgumentCaptor.forClass(Storage.BlobSourceOption.class);
    verify(storage).get(blobIdArgumentCaptor.capture());
    verify(storage).reader(any(BlobId.class), blobSourceOptionArgumentCaptor.capture());

    final BlobId blobId = blobIdArgumentCaptor.getValue();

    assertEquals(
            BUCKET,
            blobId.getBucket()
    );

    assertEquals(
            KEY,
            blobId.getName()
    );

    assertEquals(
            GENERATION,
            blobId.getGeneration()
    );


    final Set<Storage.BlobSourceOption> blobSourceOptions = ImmutableSet.copyOf(blobSourceOptionArgumentCaptor.getAllValues());
    assertTrue(blobSourceOptions.contains(Storage.BlobSourceOption.generationMatch()));
    assertEquals(
            1,
            blobSourceOptions.size()
    );

}
 

@Test
public void testKeyValueLookupAttribute() throws InitializationException {
    final SimpleKeyValueLookupService service = new SimpleKeyValueLookupService();

    final TestRunner runner = TestRunners.newTestRunner(new LookupAttribute());
    runner.addControllerService("simple-key-value-lookup-service", service);
    runner.setProperty(service, "key1", "value1");
    runner.setProperty(service, "key2", "value2");
    runner.setProperty(service, "key3", "value3");
    runner.setProperty(service, "key4", "  ");
    runner.enableControllerService(service);
    runner.assertValid(service);
    runner.setProperty(LookupAttribute.LOOKUP_SERVICE, "simple-key-value-lookup-service");
    runner.setProperty(LookupAttribute.INCLUDE_EMPTY_VALUES, "true");
    runner.setProperty("foo", "key1");
    runner.setProperty("bar", "key2");
    runner.setProperty("baz", "${attr1}");
    runner.setProperty("qux", "key4");
    runner.setProperty("zab", "key5");
    runner.assertValid();

    final Map<String, String> attributes = new HashMap<>();
    attributes.put("attr1", "key3");

    runner.enqueue("some content".getBytes(), attributes);
    runner.run(1, false);
    runner.assertAllFlowFilesTransferred(LookupAttribute.REL_MATCHED, 1);

    final MockFlowFile flowFile = runner.getFlowFilesForRelationship(LookupAttribute.REL_MATCHED).get(0);

    assertNotNull(flowFile);

    flowFile.assertAttributeExists("foo");
    flowFile.assertAttributeExists("bar");
    flowFile.assertAttributeExists("baz");
    flowFile.assertAttributeExists("qux");
    flowFile.assertAttributeExists("zab");
    flowFile.assertAttributeNotExists("zar");

    flowFile.assertAttributeEquals("foo", "value1");
    flowFile.assertAttributeEquals("bar", "value2");
    flowFile.assertAttributeEquals("baz", "value3");
    flowFile.assertAttributeEquals("qux", "");
    flowFile.assertAttributeEquals("zab", "null");
}
 

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

    final Blob blob = buildMockBlob("test-bucket-1", "test-key-1", 2L);
    final Acl.Group mockGroup = mock(Acl.Group.class);
    when(mockGroup.getEmail()).thenReturn(OWNER_GROUP_EMAIL);
    when(blob.getOwner()).thenReturn(mockGroup);

    final Iterable<Blob> mockList = ImmutableList.of(blob);

    when(mockBlobPages.getValues())
            .thenReturn(mockList);

    when(mockBlobPages.getNextPage()).thenReturn(null);

    when(storage.list(anyString(), any(Storage.BlobListOption[].class)))
            .thenReturn(mockBlobPages);

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


    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 testSuccessfulPutOperationWithUserMetadata() throws Exception {
    reset(storage, blob);
    final PutGCSObject processor = getProcessor();
    final TestRunner runner = buildNewRunner(processor);
    addRequiredPropertiesToRunner(runner);

    runner.setProperty(
            "testMetadataKey1", "testMetadataValue1"
    );
    runner.setProperty(
            "testMetadataKey2", "testMetadataValue2"
    );

    runner.assertValid();

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

    runner.enqueue("test");
    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();
    final Map<String, String> metadata = blobInfo.getMetadata();

    assertNotNull(metadata);

    assertEquals(
            2,
            metadata.size()
    );

    assertEquals(
            "testMetadataValue1",
            metadata.get("testMetadataKey1")
    );

    assertEquals(
            "testMetadataValue2",
            metadata.get("testMetadataKey2")
    );
}