Examples with PTransformOverride org.apache.beam.sdk.runners.PTransformOverride used on opensource projects

Example 6 with PTransformOverride

use of org.apache.beam.sdk.runners.PTransformOverride in project beam by apache.

the class FlinkPipelineExecutionEnvironmentTest method shouldProvideParallelismToTransformOverrides.

@Test
public void shouldProvideParallelismToTransformOverrides() {
    FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
    options.setStreaming(true);
    options.setRunner(FlinkRunner.class);
    FlinkPipelineExecutionEnvironment flinkEnv = new FlinkPipelineExecutionEnvironment(options);
    Pipeline p = Pipeline.create(options);
    // Create a transform applicable for PTransformMatchers.writeWithRunnerDeterminedSharding()
    // which requires parallelism
    p.apply(Create.of("test")).apply(TextIO.write().to("/tmp"));
    p = Mockito.spy(p);
    // If this succeeds we're ok
    flinkEnv.translate(p);
    // Verify we were using desired replacement transform
    ArgumentCaptor<ImmutableList> captor = ArgumentCaptor.forClass(ImmutableList.class);
    Mockito.verify(p).replaceAll(captor.capture());
    ImmutableList<PTransformOverride> overridesList = captor.getValue();
    assertThat(overridesList, hasItem(new BaseMatcher<PTransformOverride>() {

        @Override
        public void describeTo(Description description) {
        }

        @Override
        public boolean matches(Object actual) {
            if (actual instanceof PTransformOverride) {
                PTransformOverrideFactory overrideFactory = ((PTransformOverride) actual).getOverrideFactory();
                if (overrideFactory instanceof FlinkStreamingPipelineTranslator.StreamingShardedWriteFactory) {
                    FlinkStreamingPipelineTranslator.StreamingShardedWriteFactory factory = (FlinkStreamingPipelineTranslator.StreamingShardedWriteFactory) overrideFactory;
                    return factory.options.getParallelism() > 0;
                }
            }
            return false;
        }
    }));
}