Search in sources :

Example 6 with TestPipeline

use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.

the class PubsubIOTest method testRuntimeValueProviderSubscription.

@Test
public void testRuntimeValueProviderSubscription() {
    TestPipeline pipeline = TestPipeline.create();
    ValueProvider<String> subscription = pipeline.newProvider("projects/project/subscriptions/subscription");
    Read<String> pubsubRead = PubsubIO.readStrings().fromSubscription(subscription);
    pipeline.apply(pubsubRead);
    assertThat(pubsubRead.getSubscriptionProvider(), not(nullValue()));
    assertThat(pubsubRead.getSubscriptionProvider().isAccessible(), is(false));
}
Also used : TestPipeline(org.apache.beam.sdk.testing.TestPipeline) ByteString(com.google.protobuf.ByteString) Test(org.junit.Test)

Example 7 with TestPipeline

use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.

the class DirectGroupByKeyOverrideFactoryTest method getInputSucceeds.

@Test
public void getInputSucceeds() {
    TestPipeline p = TestPipeline.create();
    PCollection<KV<String, Integer>> input = p.apply(Create.of(KV.of("foo", 1)).withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));
    PCollection<KV<String, Iterable<Integer>>> grouped = input.apply(GroupByKey.create());
    AppliedPTransform<?, ?, ?> producer = DirectGraphs.getProducer(grouped);
    PTransformReplacement<PCollection<KV<String, Integer>>, PCollection<KV<String, Iterable<Integer>>>> replacement = factory.getReplacementTransform((AppliedPTransform) producer);
    assertThat(replacement.getInput(), Matchers.<PCollection<?>>equalTo(input));
}
Also used : PCollection(org.apache.beam.sdk.values.PCollection) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) KV(org.apache.beam.sdk.values.KV) Test(org.junit.Test)

Example 8 with TestPipeline

use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.

the class DataflowPTransformMatchersTest method createCombineGroupedValuesPipeline.

/**
 * Creates a simple pipeline with a {@link Combine.GroupedValues}.
 */
private static TestPipeline createCombineGroupedValuesPipeline() {
    TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
    PCollection<KV<String, Integer>> input = pipeline.apply(Create.of(KV.of("key", 1))).setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
    input.apply(GroupByKey.create()).apply(Combine.groupedValues(new SumCombineFn()));
    return pipeline;
}
Also used : TestPipeline(org.apache.beam.sdk.testing.TestPipeline) KV(org.apache.beam.sdk.values.KV)

Example 9 with TestPipeline

use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.

the class DataflowPTransformMatchersTest method createCombinePerKeyPipeline.

/**
 * Creates a simple pipeline with a {@link Combine.PerKey}.
 */
private static TestPipeline createCombinePerKeyPipeline() {
    TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
    PCollection<KV<String, Integer>> input = pipeline.apply(Create.of(KV.of("key", 1))).setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
    input.apply(Combine.perKey(new SumCombineFn()));
    return pipeline;
}
Also used : TestPipeline(org.apache.beam.sdk.testing.TestPipeline) KV(org.apache.beam.sdk.values.KV)

Example 10 with TestPipeline

use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.

the class DataflowPTransformMatchersTest method createCombineGroupedValuesWithSideInputsPipeline.

/**
 * Creates a simple pipeline with a {@link Combine.GroupedValues} with side inputs.
 */
private static TestPipeline createCombineGroupedValuesWithSideInputsPipeline() {
    TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
    PCollection<KV<String, Integer>> input = pipeline.apply(Create.of(KV.of("key", 1))).setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
    PCollection<String> sideInput = pipeline.apply(Create.of("side input"));
    PCollectionView<String> sideInputView = sideInput.apply(View.asSingleton());
    input.apply(GroupByKey.create()).apply(Combine.<String, Integer, Integer>groupedValues(new SumCombineFnWithContext()).withSideInputs(sideInputView));
    return pipeline;
}
Also used : TestPipeline(org.apache.beam.sdk.testing.TestPipeline) KV(org.apache.beam.sdk.values.KV)

Aggregations

TestPipeline (org.apache.beam.sdk.testing.TestPipeline)19 Test (org.junit.Test)11 KV (org.apache.beam.sdk.values.KV)5 PCollection (org.apache.beam.sdk.values.PCollection)3 ByteString (com.google.protobuf.ByteString)2 Pipeline (org.apache.beam.sdk.Pipeline)2 WriteFiles (org.apache.beam.sdk.io.WriteFiles)2 WriteFilesResult (org.apache.beam.sdk.io.WriteFilesResult)2 StorageObject (com.google.api.services.storage.model.StorageObject)1 EqualsTester (com.google.common.testing.EqualsTester)1 StreamingShardedWriteFactory (org.apache.beam.runners.dataflow.DataflowRunner.StreamingShardedWriteFactory)1 StreamingShardedWriteFactory (org.apache.beam.runners.flink.FlinkStreamingPipelineTranslator.StreamingShardedWriteFactory)1 EuphoriaOptions (org.apache.beam.sdk.extensions.euphoria.core.translate.EuphoriaOptions)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 ReplacementOutput (org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput)1 Instant (org.joda.time.Instant)1 Before (org.junit.Before)1 Parameters (org.junit.runners.Parameterized.Parameters)1 AssertResultProperties (org.talend.components.adapter.beam.example.AssertResultProperties)1 AssertResultSink (org.talend.components.adapter.beam.example.AssertResultSink)1