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

Example 16 with TestRunner

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

the class ITPutS3Object method testStatePersistsETags.

@Test
public void testStatePersistsETags() throws IOException {
    final PutS3Object processor = new PutS3Object();
    final TestRunner runner = TestRunners.newTestRunner(processor);
    final String bucket = runner.getProcessContext().getProperty(PutS3Object.BUCKET).getValue();
    final String key = runner.getProcessContext().getProperty(PutS3Object.KEY).getValue();
    final String cacheKey1 = runner.getProcessor().getIdentifier() + "/" + bucket + "/" + key + "-bv1";
    final String cacheKey2 = runner.getProcessor().getIdentifier() + "/" + bucket + "/" + key + "-bv2";
    final String cacheKey3 = runner.getProcessor().getIdentifier() + "/" + bucket + "/" + key + "-bv3";
    /*
         * store 3 versions of state
         */
    PutS3Object.MultipartState state1orig = new PutS3Object.MultipartState();
    processor.persistLocalState(cacheKey1, state1orig);
    PutS3Object.MultipartState state2orig = new PutS3Object.MultipartState();
    state2orig.setUploadId("1234");
    state2orig.setContentLength(1234L);
    processor.persistLocalState(cacheKey2, state2orig);
    PutS3Object.MultipartState state3orig = new PutS3Object.MultipartState();
    state3orig.setUploadId("5678");
    state3orig.setContentLength(5678L);
    processor.persistLocalState(cacheKey3, state3orig);
    /*
         * persist state to caches so that
         *      1. v2 has 2 and then 4 tags
         *      2. v3 has 4 and then 2 tags
         */
    state2orig.getPartETags().add(new PartETag(1, "state 2 tag one"));
    state2orig.getPartETags().add(new PartETag(2, "state 2 tag two"));
    processor.persistLocalState(cacheKey2, state2orig);
    state2orig.getPartETags().add(new PartETag(3, "state 2 tag three"));
    state2orig.getPartETags().add(new PartETag(4, "state 2 tag four"));
    processor.persistLocalState(cacheKey2, state2orig);
    state3orig.getPartETags().add(new PartETag(1, "state 3 tag one"));
    state3orig.getPartETags().add(new PartETag(2, "state 3 tag two"));
    state3orig.getPartETags().add(new PartETag(3, "state 3 tag three"));
    state3orig.getPartETags().add(new PartETag(4, "state 3 tag four"));
    processor.persistLocalState(cacheKey3, state3orig);
    state3orig.getPartETags().remove(state3orig.getPartETags().size() - 1);
    state3orig.getPartETags().remove(state3orig.getPartETags().size() - 1);
    processor.persistLocalState(cacheKey3, state3orig);
    final List<MultipartUpload> uploadList = new ArrayList<>();
    final MultipartUpload upload1 = new MultipartUpload();
    upload1.setKey(key + "-bv2");
    upload1.setUploadId("1234");
    uploadList.add(upload1);
    final MultipartUpload upload2 = new MultipartUpload();
    upload2.setKey(key + "-bv3");
    upload2.setUploadId("5678");
    uploadList.add(upload2);
    final MultipartUploadListing uploadListing = new MultipartUploadListing();
    uploadListing.setMultipartUploads(uploadList);
    final MockAmazonS3Client mockClient = new MockAmazonS3Client();
    mockClient.setListing(uploadListing);
    /*
         * load state and validate that
         *     1. v2 restore shows 4 tags
         *     2. v3 restore shows 2 tags
         */
    final PutS3Object.MultipartState state2new = processor.getLocalStateIfInS3(mockClient, bucket, cacheKey2);
    Assert.assertEquals("1234", state2new.getUploadId());
    Assert.assertEquals(4, state2new.getPartETags().size());
    final PutS3Object.MultipartState state3new = processor.getLocalStateIfInS3(mockClient, bucket, cacheKey3);
    Assert.assertEquals("5678", state3new.getUploadId());
    Assert.assertEquals(2, state3new.getPartETags().size());
}

Example 17 with TestRunner

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

the class ITPutS3Object method testMultipartSmallerThanMinimum.

@Test
public void testMultipartSmallerThanMinimum() throws IOException {
    final String FILE1_NAME = "file1";
    final byte[] megabyte = new byte[1024 * 1024];
    final Path tempFile = Files.createTempFile("s3mulitpart", "tmp");
    final FileOutputStream tempOut = new FileOutputStream(tempFile.toFile());
    long tempByteCount = 0;
    for (int i = 0; i < 5; i++) {
        tempOut.write(megabyte);
        tempByteCount += megabyte.length;
    }
    tempOut.close();
    System.out.println("file size: " + tempByteCount);
    Assert.assertTrue(tempByteCount < S3_MINIMUM_PART_SIZE);
    Assert.assertTrue(megabyte.length < S3_MINIMUM_PART_SIZE);
    Assert.assertTrue(TEST_PARTSIZE_LONG >= S3_MINIMUM_PART_SIZE && TEST_PARTSIZE_LONG <= S3_MAXIMUM_OBJECT_SIZE);
    final PutS3Object processor = new PutS3Object();
    final TestRunner runner = TestRunners.newTestRunner(processor);
    runner.setProperty(PutS3Object.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(PutS3Object.REGION, REGION);
    runner.setProperty(PutS3Object.BUCKET, BUCKET_NAME);
    runner.setProperty(PutS3Object.MULTIPART_PART_SIZE, TEST_PARTSIZE_STRING);
    Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.FILENAME.key(), FILE1_NAME);
    runner.enqueue(new FileInputStream(tempFile.toFile()), attributes);
    runner.assertValid();
    runner.run();
    runner.assertAllFlowFilesTransferred(PutS3Object.REL_SUCCESS);
    final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(PutS3Object.REL_SUCCESS);
    Assert.assertEquals(1, successFiles.size());
    final List<MockFlowFile> failureFiles = runner.getFlowFilesForRelationship(PutS3Object.REL_FAILURE);
    Assert.assertEquals(0, failureFiles.size());
    MockFlowFile ff1 = successFiles.get(0);
    Assert.assertEquals(PutS3Object.S3_API_METHOD_PUTOBJECT, ff1.getAttribute(PutS3Object.S3_API_METHOD_ATTR_KEY));
    Assert.assertEquals(FILE1_NAME, ff1.getAttribute(CoreAttributes.FILENAME.key()));
    Assert.assertEquals(BUCKET_NAME, ff1.getAttribute(PutS3Object.S3_BUCKET_KEY));
    Assert.assertEquals(FILE1_NAME, ff1.getAttribute(PutS3Object.S3_OBJECT_KEY));
    Assert.assertTrue(reS3ETag.matcher(ff1.getAttribute(PutS3Object.S3_ETAG_ATTR_KEY)).matches());
    Assert.assertEquals(tempByteCount, ff1.getSize());
}

Example 18 with TestRunner

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

the class ITPutS3Object method testEndpointOverride.

@Test
public void testEndpointOverride() {
    // remove leading "/" from filename to avoid duplicate separators
    final String TESTKEY = AbstractS3IT.SAMPLE_FILE_RESOURCE_NAME.substring(1);
    final PutS3Object processor = new TestablePutS3Object();
    final TestRunner runner = TestRunners.newTestRunner(processor);
    final ProcessContext context = runner.getProcessContext();
    runner.setProperty(PutS3Object.ENDPOINT_OVERRIDE, TEST_ENDPOINT);
    runner.setProperty(PutS3Object.BUCKET, BUCKET_NAME);
    runner.setProperty(PutS3Object.KEY, TESTKEY);
    runner.run();
    Assert.assertEquals(BUCKET_NAME, context.getProperty(PutS3Object.BUCKET).toString());
    Assert.assertEquals(TESTKEY, context.getProperty(PutS3Object.KEY).evaluateAttributeExpressions().toString());
    Assert.assertEquals(TEST_ENDPOINT, context.getProperty(PutS3Object.ENDPOINT_OVERRIDE).toString());
    String s3url = ((TestablePutS3Object) processor).testable_getClient().getResourceUrl(BUCKET_NAME, TESTKEY);
    Assert.assertEquals(TEST_ENDPOINT + "/" + BUCKET_NAME + "/" + TESTKEY, s3url);
}

Example 19 with TestRunner

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

the class ITPutS3Object method testPutThenFetch.

private void testPutThenFetch(String sseAlgorithm) throws IOException {
    // Put
    TestRunner runner = TestRunners.newTestRunner(new PutS3Object());
    runner.setProperty(PutS3Object.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(PutS3Object.REGION, REGION);
    runner.setProperty(PutS3Object.BUCKET, BUCKET_NAME);
    if (ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION.equals(sseAlgorithm)) {
        runner.setProperty(PutS3Object.SERVER_SIDE_ENCRYPTION, sseAlgorithm);
    }
    final Map<String, String> attrs = new HashMap<>();
    attrs.put("filename", "filename-on-s3.txt");
    runner.enqueue(getResourcePath(SAMPLE_FILE_RESOURCE_NAME), attrs);
    runner.run();
    runner.assertAllFlowFilesTransferred(PutS3Object.REL_SUCCESS, 1);
    List<MockFlowFile> ffs = runner.getFlowFilesForRelationship(PutS3Object.REL_SUCCESS);
    if (ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION.equals(sseAlgorithm)) {
        ffs.get(0).assertAttributeEquals(PutS3Object.S3_SSE_ALGORITHM, ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION);
    } else {
        ffs.get(0).assertAttributeNotExists(PutS3Object.S3_SSE_ALGORITHM);
    }
    // Fetch
    runner = TestRunners.newTestRunner(new FetchS3Object());
    runner.setProperty(FetchS3Object.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(FetchS3Object.REGION, REGION);
    runner.setProperty(FetchS3Object.BUCKET, BUCKET_NAME);
    runner.enqueue(new byte[0], attrs);
    runner.run(1);
    runner.assertAllFlowFilesTransferred(FetchS3Object.REL_SUCCESS, 1);
    ffs = runner.getFlowFilesForRelationship(FetchS3Object.REL_SUCCESS);
    MockFlowFile ff = ffs.get(0);
    ff.assertContentEquals(getFileFromResourceName(SAMPLE_FILE_RESOURCE_NAME));
    if (ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION.equals(sseAlgorithm)) {
        ff.assertAttributeEquals(PutS3Object.S3_SSE_ALGORITHM, ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION);
    } else {
        ff.assertAttributeNotExists(PutS3Object.S3_SSE_ALGORITHM);
    }
}

Example 20 with TestRunner

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

the class ITPutS3Object method testPutInFolder.

@Test
public void testPutInFolder() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new PutS3Object());
    runner.setProperty(PutS3Object.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(PutS3Object.REGION, REGION);
    runner.setProperty(PutS3Object.BUCKET, BUCKET_NAME);
    Assert.assertTrue(runner.setProperty("x-custom-prop", "hello").isValid());
    runner.assertValid();
    final Map<String, String> attrs = new HashMap<>();
    attrs.put("filename", "folder/1.txt");
    runner.enqueue(getResourcePath(SAMPLE_FILE_RESOURCE_NAME), attrs);
    runner.run();
    runner.assertAllFlowFilesTransferred(PutS3Object.REL_SUCCESS, 1);
}