Search in sources :

Example 6 with DeterministicallyConstructTestRowFn

use of org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn in project beam by apache.

the class SqsIOIT method testWriteThenRead.

@Test
public void testWriteThenRead() {
    int rows = env.options().getNumberOfRows();
    // Write test dataset to SQS.
    pipelineWrite.apply("Generate Sequence", GenerateSequence.from(0).to(rows)).apply("Prepare TestRows", ParDo.of(new DeterministicallyConstructTestRowFn())).apply("Prepare SQS message", MapElements.into(requestType).via(sqsQueue::messageRequest)).apply("Write to SQS", SqsIO.write());
    // Read test dataset from SQS.
    PCollection<String> output = pipelineRead.apply("Read from SQS", SqsIO.read().withQueueUrl(sqsQueue.url).withMaxNumRecords(rows)).apply("Extract body", MapElements.into(strings()).via(SqsMessage::getBody));
    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 : DeterministicallyConstructTestRowFn(org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn) HashingFn(org.apache.beam.sdk.io.common.HashingFn) Test(org.junit.Test)

Example 7 with DeterministicallyConstructTestRowFn

use of org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn 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("SNS request", MapElements.into(publishRequests).via(resources::publishRequest)).apply("Write to SNS", SnsIO.write().withTopicName(resources.snsTopic).withResultOutputTag(new TupleTag<>()).withAWSClientsProvider(opts.getAwsCredentialsProvider().getCredentials().getAWSAccessKeyId(), opts.getAwsCredentialsProvider().getCredentials().getAWSSecretKey(), Regions.fromName(opts.getAwsRegion()), opts.getAwsServiceEndpoint()));
    // 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 : DeterministicallyConstructTestRowFn(org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn) HashingFn(org.apache.beam.sdk.io.common.HashingFn) Test(org.junit.Test)

Example 8 with DeterministicallyConstructTestRowFn

use of org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn in project beam by apache.

the class SqsIOIT method testWriteThenRead.

@Test
public void testWriteThenRead() {
    int rows = env.options().getNumberOfRows();
    // Write test dataset to SQS.
    pipelineWrite.apply("Generate Sequence", GenerateSequence.from(0).to(rows)).apply("Prepare TestRows", ParDo.of(new DeterministicallyConstructTestRowFn())).apply("Prepare SQS message", MapElements.into(requestType).via(sqsQueue::messageRequest)).apply("Write to SQS", SqsIO.write());
    // Read test dataset from SQS.
    PCollection<String> output = pipelineRead.apply("Read from SQS", SqsIO.read().withQueueUrl(sqsQueue.url).withMaxNumRecords(rows)).apply("Extract body", MapElements.into(strings()).via(Message::getBody));
    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 : DeterministicallyConstructTestRowFn(org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn) HashingFn(org.apache.beam.sdk.io.common.HashingFn) Test(org.junit.Test)

Aggregations

HashingFn (org.apache.beam.sdk.io.common.HashingFn)8 DeterministicallyConstructTestRowFn (org.apache.beam.sdk.io.common.TestRow.DeterministicallyConstructTestRowFn)8 Test (org.junit.Test)8 GenerateSequence (org.apache.beam.sdk.io.GenerateSequence)3 TestRow (org.apache.beam.sdk.io.common.TestRow)3 TestRow.getExpectedHashForRowCount (org.apache.beam.sdk.io.common.TestRow.getExpectedHashForRowCount)3 PAssert (org.apache.beam.sdk.testing.PAssert)3 TestPipeline (org.apache.beam.sdk.testing.TestPipeline)3 Combine (org.apache.beam.sdk.transforms.Combine)3 Count (org.apache.beam.sdk.transforms.Count)3 MapElements (org.apache.beam.sdk.transforms.MapElements)3 ParDo (org.apache.beam.sdk.transforms.ParDo)3 PCollection (org.apache.beam.sdk.values.PCollection)3 TypeDescriptors.strings (org.apache.beam.sdk.values.TypeDescriptors.strings)3 ClassRule (org.junit.ClassRule)3 Rule (org.junit.Rule)3 ExternalResource (org.junit.rules.ExternalResource)3 RunWith (org.junit.runner.RunWith)3 JUnit4 (org.junit.runners.JUnit4)3 Map (java.util.Map)2