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

Example 31 with MockFlowFile

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

the class ITPutKinesisStream method testIntegrationWithFixedPartitionSuccess.

@Test
public void testIntegrationWithFixedPartitionSuccess() throws Exception {
    runner.setProperty(PutKinesisStream.KINESIS_PARTITION_KEY, "pfixed");
    runner.assertValid();
    runner.enqueue("test".getBytes());
    runner.run(1);
    runner.assertAllFlowFilesTransferred(PutKinesisStream.REL_SUCCESS, 1);
    final List<MockFlowFile> ffs = runner.getFlowFilesForRelationship(PutKinesisStream.REL_SUCCESS);
    final MockFlowFile out = ffs.iterator().next();
    out.assertContentEquals("test".getBytes());
}

Example 32 with MockFlowFile

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

the class ITPutKinesisStream method testThreeMessageWithBatch5MaxBufferSize1MBRunOnceTwoMessageSent.

@Test
public void testThreeMessageWithBatch5MaxBufferSize1MBRunOnceTwoMessageSent() {
    runner.setProperty(PutKinesisStream.BATCH_SIZE, "5");
    runner.setProperty(PutKinesisStream.MAX_MESSAGE_BUFFER_SIZE_MB, "1 MB");
    runner.assertValid();
    byte[] bytes = new byte[(PutKinesisStream.MAX_MESSAGE_SIZE)];
    for (int i = 0; i < bytes.length; i++) {
        bytes[i] = 'a';
    }
    runner.enqueue(bytes);
    runner.enqueue(bytes.clone());
    runner.enqueue(bytes.clone());
    runner.run(1, true, true);
    runner.assertAllFlowFilesTransferred(PutKinesisStream.REL_SUCCESS, 2);
    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(PutKinesisStream.REL_SUCCESS);
    assertEquals(2, flowFiles.size());
    for (MockFlowFile flowFile : flowFiles) {
        flowFile.assertAttributeExists(PutKinesisStream.AWS_KINESIS_SEQUENCE_NUMBER);
        flowFile.assertAttributeExists(PutKinesisStream.AWS_KINESIS_SHARD_ID);
    }
}

Example 33 with MockFlowFile

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

the class ITPutKinesisStream method testTwoMessageBatch5WithMaxBufferSize1MBRunOnceTwoMessageSent.

@Test
public void testTwoMessageBatch5WithMaxBufferSize1MBRunOnceTwoMessageSent() {
    runner.setProperty(PutKinesisStream.BATCH_SIZE, "5");
    runner.setProperty(PutKinesisStream.MAX_MESSAGE_BUFFER_SIZE_MB, "2 MB");
    runner.assertValid();
    byte[] bytes = new byte[(PutKinesisStream.MAX_MESSAGE_SIZE)];
    for (int i = 0; i < bytes.length; i++) {
        bytes[i] = 'a';
    }
    runner.enqueue(bytes);
    runner.enqueue(bytes.clone());
    runner.run(1);
    runner.assertAllFlowFilesTransferred(PutKinesisStream.REL_SUCCESS, 2);
    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(PutKinesisStream.REL_SUCCESS);
    assertEquals(2, flowFiles.size());
    for (MockFlowFile flowFile : flowFiles) {
        flowFile.assertAttributeExists(PutKinesisStream.AWS_KINESIS_SEQUENCE_NUMBER);
        flowFile.assertAttributeExists(PutKinesisStream.AWS_KINESIS_SHARD_ID);
    }
}

Example 34 with MockFlowFile

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

the class ITPutLambda method testIntegrationSuccess.

/**
 * Comment out ignore for integration tests (requires creds files)
 */
@Test
@Ignore
public void testIntegrationSuccess() throws Exception {
    runner = TestRunners.newTestRunner(PutLambda.class);
    runner.setProperty(PutLambda.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(PutLambda.AWS_LAMBDA_FUNCTION_NAME, "hello");
    runner.assertValid();
    runner.enqueue("{\"test\":\"hi\"}".getBytes());
    runner.run(1);
    runner.assertAllFlowFilesTransferred(PutLambda.REL_SUCCESS, 1);
    final List<MockFlowFile> ffs = runner.getFlowFilesForRelationship(PutLambda.REL_SUCCESS);
    final MockFlowFile out = ffs.iterator().next();
    assertNull("Function error should be null " + out.getAttribute(PutLambda.AWS_LAMBDA_RESULT_FUNCTION_ERROR), out.getAttribute(PutLambda.AWS_LAMBDA_RESULT_FUNCTION_ERROR));
    assertNotNull("log should not be null", out.getAttribute(PutLambda.AWS_LAMBDA_RESULT_LOG));
    assertEquals("Status should be equal", "200", out.getAttribute(PutLambda.AWS_LAMBDA_RESULT_STATUS_CODE));
}

Example 35 with MockFlowFile

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

the class TestPutLambda method testPutLambdaTooManyRequestsException.

@Test
public void testPutLambdaTooManyRequestsException() {
    runner.setProperty(PutLambda.AWS_LAMBDA_FUNCTION_NAME, "test-function");
    runner.enqueue("TestContent");
    Mockito.when(mockLambdaClient.invoke(Mockito.any(InvokeRequest.class))).thenThrow(new TooManyRequestsException("TestFail"));
    runner.assertValid();
    runner.run(1);
    runner.assertAllFlowFilesTransferred(PutLambda.REL_FAILURE, 1);
    final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(PutLambda.REL_FAILURE);
    final MockFlowFile ff0 = flowFiles.get(0);
    assertTrue(ff0.isPenalized());
}