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

Example 21 with TestRunner

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

the class ITPutS3Object method testLocalStatePersistence.

@Test
public void testLocalStatePersistence() 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;
    final String cacheKey2 = runner.getProcessor().getIdentifier() + "/" + bucket + "/" + key + "-v2";
    final String cacheKey3 = runner.getProcessor().getIdentifier() + "/" + bucket + "/" + key + "-v3";
    /*
         * 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);
    final List<MultipartUpload> uploadList = new ArrayList<>();
    final MultipartUpload upload1 = new MultipartUpload();
    upload1.setKey(key);
    upload1.setUploadId("");
    uploadList.add(upload1);
    final MultipartUpload upload2 = new MultipartUpload();
    upload2.setKey(key + "-v2");
    upload2.setUploadId("1234");
    uploadList.add(upload2);
    final MultipartUpload upload3 = new MultipartUpload();
    upload3.setKey(key + "-v3");
    upload3.setUploadId("5678");
    uploadList.add(upload3);
    final MultipartUploadListing uploadListing = new MultipartUploadListing();
    uploadListing.setMultipartUploads(uploadList);
    final MockAmazonS3Client mockClient = new MockAmazonS3Client();
    mockClient.setListing(uploadListing);
    /*
         * reload and validate stored state
         */
    final PutS3Object.MultipartState state1new = processor.getLocalStateIfInS3(mockClient, bucket, cacheKey1);
    Assert.assertEquals("", state1new.getUploadId());
    Assert.assertEquals(0L, state1new.getFilePosition().longValue());
    Assert.assertEquals(new ArrayList<PartETag>(), state1new.getPartETags());
    Assert.assertEquals(0L, state1new.getPartSize().longValue());
    Assert.assertEquals(StorageClass.fromValue(StorageClass.Standard.toString()), state1new.getStorageClass());
    Assert.assertEquals(0L, state1new.getContentLength().longValue());
    final PutS3Object.MultipartState state2new = processor.getLocalStateIfInS3(mockClient, bucket, cacheKey2);
    Assert.assertEquals("1234", state2new.getUploadId());
    Assert.assertEquals(0L, state2new.getFilePosition().longValue());
    Assert.assertEquals(new ArrayList<PartETag>(), state2new.getPartETags());
    Assert.assertEquals(0L, state2new.getPartSize().longValue());
    Assert.assertEquals(StorageClass.fromValue(StorageClass.Standard.toString()), state2new.getStorageClass());
    Assert.assertEquals(1234L, state2new.getContentLength().longValue());
    final PutS3Object.MultipartState state3new = processor.getLocalStateIfInS3(mockClient, bucket, cacheKey3);
    Assert.assertEquals("5678", state3new.getUploadId());
    Assert.assertEquals(0L, state3new.getFilePosition().longValue());
    Assert.assertEquals(new ArrayList<PartETag>(), state3new.getPartETags());
    Assert.assertEquals(0L, state3new.getPartSize().longValue());
    Assert.assertEquals(StorageClass.fromValue(StorageClass.Standard.toString()), state3new.getStorageClass());
    Assert.assertEquals(5678L, state3new.getContentLength().longValue());
}

Example 22 with TestRunner

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

the class ITPutS3Object method testDynamicProperty.

@Test
public void testDynamicProperty() throws IOException {
    final String DYNAMIC_ATTRIB_KEY = "fs.runTimestamp";
    final String DYNAMIC_ATTRIB_VALUE = "${now():toNumber()}";
    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);
    PropertyDescriptor testAttrib = processor.getSupportedDynamicPropertyDescriptor(DYNAMIC_ATTRIB_KEY);
    runner.setProperty(testAttrib, DYNAMIC_ATTRIB_VALUE);
    final String FILE1_NAME = "file1";
    Map<String, String> attribs = new HashMap<>();
    attribs.put(CoreAttributes.FILENAME.key(), FILE1_NAME);
    runner.enqueue("123".getBytes(), attribs);
    runner.assertValid();
    processor.getPropertyDescriptor(DYNAMIC_ATTRIB_KEY);
    runner.run();
    runner.assertAllFlowFilesTransferred(PutS3Object.REL_SUCCESS);
    final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(PutS3Object.REL_SUCCESS);
    Assert.assertEquals(1, successFiles.size());
    MockFlowFile ff1 = successFiles.get(0);
    Long now = System.currentTimeMillis();
    String millisNow = Long.toString(now);
    String millisOneSecAgo = Long.toString(now - 1000L);
    String usermeta = ff1.getAttribute(PutS3Object.S3_USERMETA_ATTR_KEY);
    String[] usermetaLine0 = usermeta.split(System.lineSeparator())[0].split("=");
    String usermetaKey0 = usermetaLine0[0];
    String usermetaValue0 = usermetaLine0[1];
    Assert.assertEquals(DYNAMIC_ATTRIB_KEY, usermetaKey0);
    Assert.assertTrue(usermetaValue0.compareTo(millisOneSecAgo) >= 0 && usermetaValue0.compareTo(millisNow) <= 0);
}

Example 23 with TestRunner

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

the class TestPutS3Object method testSignerOverrideOptions.

@Test
public void testSignerOverrideOptions() {
    final AWSCredentialsProvider credentialsProvider = new DefaultAWSCredentialsProviderChain();
    final ClientConfiguration config = new ClientConfiguration();
    final PutS3Object processor = new PutS3Object();
    final TestRunner runner = TestRunners.newTestRunner(processor);
    final List<AllowableValue> allowableSignerValues = PutS3Object.SIGNER_OVERRIDE.getAllowableValues();
    final String defaultSignerValue = PutS3Object.SIGNER_OVERRIDE.getDefaultValue();
    for (AllowableValue allowableSignerValue : allowableSignerValues) {
        String signerType = allowableSignerValue.getValue();
        if (!signerType.equals(defaultSignerValue)) {
            runner.setProperty(PutS3Object.SIGNER_OVERRIDE, signerType);
            ProcessContext context = runner.getProcessContext();
            try {
                AmazonS3Client s3Client = processor.createClient(context, credentialsProvider, config);
            } catch (IllegalArgumentException argEx) {
                Assert.fail(argEx.getMessage());
            }
        }
    }
}

Example 24 with TestRunner

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

the class ITPutSNS method testPublish.

@Test
public void testPublish() throws IOException {
    final TestRunner runner = TestRunners.newTestRunner(new PutSNS());
    runner.setValidateExpressionUsage(false);
    runner.setProperty(PutSNS.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(PutSNS.ARN, "arn:aws:sns:us-west-2:100515378163:test-topic-1");
    assertTrue(runner.setProperty("DynamicProperty", "hello!").isValid());
    final Map<String, String> attrs = new HashMap<>();
    attrs.put("filename", "1.txt");
    runner.enqueue(Paths.get("src/test/resources/hello.txt"), attrs);
    runner.run();
    runner.assertAllFlowFilesTransferred(PutSNS.REL_SUCCESS, 1);
}

Example 25 with TestRunner

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

the class ITGetSQS method testSimpleGetUsingCredentialsProviderService.

@Test
public void testSimpleGetUsingCredentialsProviderService() throws Throwable {
    final TestRunner runner = TestRunners.newTestRunner(new GetSQS());
    runner.setProperty(GetSQS.TIMEOUT, "30 secs");
    String queueUrl = "Add queue url here";
    runner.setProperty(GetSQS.QUEUE_URL, queueUrl);
    final AWSCredentialsProviderControllerService serviceImpl = new AWSCredentialsProviderControllerService();
    runner.addControllerService("awsCredentialsProvider", serviceImpl);
    runner.setProperty(serviceImpl, AbstractAWSProcessor.CREDENTIALS_FILE, System.getProperty("user.home") + "/aws-credentials.properties");
    runner.enableControllerService(serviceImpl);
    runner.assertValid(serviceImpl);
    runner.setProperty(GetSQS.AWS_CREDENTIALS_PROVIDER_SERVICE, "awsCredentialsProvider");
    runner.run(1);
    final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(GetSQS.REL_SUCCESS);
    for (final MockFlowFile mff : flowFiles) {
        System.out.println(mff.getAttributes());
        System.out.println(new String(mff.toByteArray()));
    }
}