use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.
the class DataflowPTransformMatchersTest method createCombinePerKeyWithSideInputsPipeline.
/**
* Creates a simple pipeline with a {@link Combine.PerKey} with side inputs.
*/
private static TestPipeline createCombinePerKeyWithSideInputsPipeline() {
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(Combine.<String, Integer, Integer>perKey(new SumCombineFnWithContext()).withSideInputs(sideInputView));
return pipeline;
}
use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.
the class PTransformTranslationTest method data.
@Parameters(name = "{index}: {0}")
public static Iterable<ToAndFromProtoSpec> data() {
// This pipeline exists for construction, not to run any test.
// TODO: Leaf node with understood payload - i.e. validate payloads
ToAndFromProtoSpec readLeaf = ToAndFromProtoSpec.leaf(read(TestPipeline.create()));
ToAndFromProtoSpec readMultipleInAndOut = ToAndFromProtoSpec.leaf(multiMultiParDo(TestPipeline.create()));
TestPipeline compositeReadPipeline = TestPipeline.create();
ToAndFromProtoSpec compositeRead = ToAndFromProtoSpec.composite(generateSequence(compositeReadPipeline), ToAndFromProtoSpec.leaf(read(compositeReadPipeline)));
ToAndFromProtoSpec rawLeafNullSpec = ToAndFromProtoSpec.leaf(rawPTransformWithNullSpec(TestPipeline.create()));
return ImmutableList.<ToAndFromProtoSpec>builder().add(readLeaf).add(readMultipleInAndOut).add(compositeRead).add(rawLeafNullSpec).build();
}
use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.
the class UnionTest method testBuild_ImplicitName.
@Test
public void testBuild_ImplicitName() {
final TestPipeline pipeline = TestUtils.createTestPipeline();
final PCollection<String> left = TestUtils.createMockDataset(pipeline, TypeDescriptors.strings());
final PCollection<String> right = TestUtils.createMockDataset(pipeline, TypeDescriptors.strings());
final PCollection<String> unioned = Union.of(left, right).output();
final Union union = (Union) TestUtils.getProducer(unioned);
assertFalse(union.getName().isPresent());
}
use of org.apache.beam.sdk.testing.TestPipeline in project beam by apache.
the class UnionTest method testBuild.
@Test
public void testBuild() {
final TestPipeline pipeline = TestUtils.createTestPipeline();
final PCollection<String> left = TestUtils.createMockDataset(pipeline, TypeDescriptors.strings());
final PCollection<String> right = TestUtils.createMockDataset(pipeline, TypeDescriptors.strings());
final PCollection<String> unioned = Union.named("Union1").of(left, right).output();
final Union union = (Union) TestUtils.getProducer(unioned);
assertTrue(union.getName().isPresent());
assertEquals("Union1", union.getName().get());
}
Aggregations