Search in sources :

Example 1 with Read

use of org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read in project beam by apache.

the class KinesisIOReadTest method testBuildWithCredentialsProvider.

@Test
public void testBuildWithCredentialsProvider() {
    Region region = Region.US_EAST_1;
    AwsCredentialsProvider credentialsProvider = DefaultCredentialsProvider.create();
    Read read = KinesisIO.read().withAWSClientsProvider(credentialsProvider, region);
    assertThat(read.getClientConfiguration()).isEqualTo(ClientConfiguration.create(credentialsProvider, region, null));
}
Also used : Read(org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read) AwsCredentialsProvider(software.amazon.awssdk.auth.credentials.AwsCredentialsProvider) Region(software.amazon.awssdk.regions.Region) Test(org.junit.Test)

Example 2 with Read

use of org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read in project beam by apache.

the class KinesisIOReadTest method readFromShards.

private void readFromShards(Function<Read, Read> fn, Iterable<Record> expected) {
    Read read = KinesisIO.read().withStreamName("stream").withInitialPositionInStream(TRIM_HORIZON).withArrivalTimeWatermarkPolicy().withMaxNumRecords(SHARDS * SHARD_EVENTS);
    PCollection<Record> result = p.apply(fn.apply(read)).apply(ParDo.of(new ToRecord()));
    PAssert.that(result).containsInAnyOrder(expected);
    p.run();
}
Also used : Read(org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read) Record(software.amazon.awssdk.services.kinesis.model.Record)

Example 3 with Read

use of org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read in project beam by apache.

the class SqsIOReadTest method testBuildWithCredentialsProviderAndRegionAndEndpoint.

@Test
public void testBuildWithCredentialsProviderAndRegionAndEndpoint() {
    Region region = Region.US_EAST_1;
    AwsCredentialsProvider credentialsProvider = DefaultCredentialsProvider.create();
    URI endpoint = URI.create("localhost:9999");
    Read read = SqsIO.read().withSqsClientProvider(credentialsProvider, region.id(), endpoint);
    assertThat(read.clientConfiguration()).isEqualTo(ClientConfiguration.create(credentialsProvider, region, endpoint));
}
Also used : Read(org.apache.beam.sdk.io.aws2.sqs.SqsIO.Read) AwsCredentialsProvider(software.amazon.awssdk.auth.credentials.AwsCredentialsProvider) Region(software.amazon.awssdk.regions.Region) URI(java.net.URI) Test(org.junit.Test)

Example 4 with Read

use of org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read in project beam by apache.

the class SnsIOIT method testWriteThenRead.

@Test
public void testWriteThenRead() {
    ITOptions opts = env.options();
    int rows = opts.getNumberOfRows();
    // Write test dataset to SNS
    pipelineWrite.apply("Generate Sequence", GenerateSequence.from(0).to(rows)).apply("Prepare TestRows", ParDo.of(new DeterministicallyConstructTestRowFn())).apply("Write to SNS", SnsIO.<TestRow>write().withTopicArn(resources.snsTopic).withPublishRequestBuilder(r -> PublishRequest.builder().message(r.name())));
    // Read test dataset from SQS.
    PCollection<String> output = pipelineRead.apply("Read from SQS", SqsIO.read().withQueueUrl(resources.sqsQueue).withMaxNumRecords(rows)).apply("Extract message", MapElements.into(strings()).via(SnsIOIT::extractMessage));
    PAssert.thatSingleton(output.apply("Count All", Count.globally())).isEqualTo((long) rows);
    PAssert.that(output.apply(Combine.globally(new HashingFn()).withoutDefaults())).containsInAnyOrder(getExpectedHashForRowCount(rows));
    pipelineWrite.run();
    pipelineRead.run();
}
Also used : Count(org.apache.beam.sdk.transforms.Count) Combine(org.apache.beam.sdk.transforms.Combine) RunWith(org.junit.runner.RunWith) PublishRequest(software.amazon.awssdk.services.sns.model.PublishRequest) IOITHelper.executeWithRetry(org.apache.beam.sdk.io.common.IOITHelper.executeWithRetry) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) TypeDescriptors.strings(org.apache.beam.sdk.values.TypeDescriptors.strings) Timeout(org.junit.rules.Timeout) TestRow.getExpectedHashForRowCount(org.apache.beam.sdk.io.common.TestRow.getExpectedHashForRowCount) ClassRule(org.junit.ClassRule) SqsIO(org.apache.beam.sdk.io.aws2.sqs.SqsIO) Service(org.testcontainers.containers.localstack.LocalStackContainer.Service) MapElements(org.apache.beam.sdk.transforms.MapElements) SNS(org.testcontainers.containers.localstack.LocalStackContainer.Service.SNS) HashingFn(org.apache.beam.sdk.io.common.HashingFn) DeterministicallyConstructTestRowFn(org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn) PAssert(org.apache.beam.sdk.testing.PAssert) TestRow(org.apache.beam.sdk.io.common.TestRow) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) ITEnvironment(org.apache.beam.sdk.io.aws2.ITEnvironment) SqsClient(software.amazon.awssdk.services.sqs.SqsClient) JsonProcessingException(com.fasterxml.jackson.core.JsonProcessingException) GenerateSequence(org.apache.beam.sdk.io.GenerateSequence) Test(org.junit.Test) JUnit4(org.junit.runners.JUnit4) PCollection(org.apache.beam.sdk.values.PCollection) Serializable(java.io.Serializable) Rule(org.junit.Rule) ExternalResource(org.junit.rules.ExternalResource) ParDo(org.apache.beam.sdk.transforms.ParDo) SnsClient(software.amazon.awssdk.services.sns.SnsClient) SqsMessage(org.apache.beam.sdk.io.aws2.sqs.SqsMessage) SQS(org.testcontainers.containers.localstack.LocalStackContainer.Service.SQS) DeterministicallyConstructTestRowFn(org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn) HashingFn(org.apache.beam.sdk.io.common.HashingFn) Test(org.junit.Test)

Example 5 with Read

use of org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read in project beam by apache.

the class S3FileSystemTest method testWriteAndRead.

@Test
public void testWriteAndRead() throws IOException {
    S3FileSystem s3FileSystem = buildMockedS3FileSystem(s3Config("mys3"), client);
    client.createBucket(CreateBucketRequest.builder().bucket("testbucket").build());
    byte[] writtenArray = new byte[] { 0 };
    ByteBuffer bb = ByteBuffer.allocate(writtenArray.length);
    bb.put(writtenArray);
    // First create an object and write data to it
    S3ResourceId path = S3ResourceId.fromUri("mys3://testbucket/foo/bar.txt");
    WritableByteChannel writableByteChannel = s3FileSystem.create(path, CreateOptions.StandardCreateOptions.builder().setMimeType("application/text").build());
    writableByteChannel.write(bb);
    writableByteChannel.close();
    // Now read the same object
    ByteBuffer bb2 = ByteBuffer.allocate(writtenArray.length);
    ReadableByteChannel open = s3FileSystem.open(path);
    open.read(bb2);
    // And compare the content with the one that was written
    byte[] readArray = bb2.array();
    assertArrayEquals(readArray, writtenArray);
    open.close();
}
Also used : ReadableByteChannel(java.nio.channels.ReadableByteChannel) WritableByteChannel(java.nio.channels.WritableByteChannel) ByteBuffer(java.nio.ByteBuffer) S3TestUtils.buildMockedS3FileSystem(org.apache.beam.sdk.io.aws2.s3.S3TestUtils.buildMockedS3FileSystem) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)11 Region (software.amazon.awssdk.regions.Region)6 Read (org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read)5 AwsCredentialsProvider (software.amazon.awssdk.auth.credentials.AwsCredentialsProvider)4 S3TestUtils.buildMockedS3FileSystem (org.apache.beam.sdk.io.aws2.s3.S3TestUtils.buildMockedS3FileSystem)3 GenerateSequence (org.apache.beam.sdk.io.GenerateSequence)2 ITEnvironment (org.apache.beam.sdk.io.aws2.ITEnvironment)2 Read (org.apache.beam.sdk.io.aws2.sqs.SqsIO.Read)2 HashingFn (org.apache.beam.sdk.io.common.HashingFn)2 TestRow (org.apache.beam.sdk.io.common.TestRow)2 DeterministicallyConstructTestRowFn (org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn)2 TestRow.getExpectedHashForRowCount (org.apache.beam.sdk.io.common.TestRow.getExpectedHashForRowCount)2 MatchResult (org.apache.beam.sdk.io.fs.MatchResult)2 PAssert (org.apache.beam.sdk.testing.PAssert)2 TestPipeline (org.apache.beam.sdk.testing.TestPipeline)2 Combine (org.apache.beam.sdk.transforms.Combine)2 Count (org.apache.beam.sdk.transforms.Count)2 MapElements (org.apache.beam.sdk.transforms.MapElements)2 ParDo (org.apache.beam.sdk.transforms.ParDo)2 PCollection (org.apache.beam.sdk.values.PCollection)2