Examples with Receiver org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver used on opensource projects

Example 6 with Receiver

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver in project beam by apache.

the class UserParDoFnFactoryTest method testFactorySimultaneousUse.

@Test
public void testFactorySimultaneousUse() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    CounterSet counters = new CounterSet();
    TestDoFn initialFn = new TestDoFn(Collections.<TupleTag<String>>emptyList());
    CloudObject cloudObject = getCloudObject(initialFn);
    ParDoFn parDoFn = factory.create(options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters));
    // The fn should not be reused while the first ParDoFn is not finished
    ParDoFn secondParDoFn = factory.create(options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters));
    Receiver rcvr = new OutputReceiver();
    parDoFn.startBundle(rcvr);
    parDoFn.processElement(WindowedValue.valueInGlobalWindow("foo"));
    // Must be after the first call to process element for reallyStartBundle to have been called
    TestDoFn firstDoFn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn();
    secondParDoFn.startBundle(rcvr);
    secondParDoFn.processElement(WindowedValue.valueInGlobalWindow("spam"));
    // Must be after the first call to process element for reallyStartBundle to have been called
    TestDoFn secondDoFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn();
    parDoFn.finishBundle();
    secondParDoFn.finishBundle();
    assertThat(firstDoFn, not(theInstance(secondDoFn)));
    assertThat(firstDoFn.state, equalTo(TestDoFn.State.FINISHED));
    assertThat(secondDoFn.state, equalTo(TestDoFn.State.FINISHED));
}