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));
}
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));
}
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;
}
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;
}
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;
}
Aggregations