Examples with HandlesSplits org.apache.beam.fn.harness.HandlesSplits used on opensource projects

Example 1 with HandlesSplits

use of org.apache.beam.fn.harness.HandlesSplits in project beam by apache.

the class PCollectionConsumerRegistryTest method testHandlesSplitsPassedToOriginalConsumer.

@Test
public void testHandlesSplitsPassedToOriginalConsumer() throws Exception {
    final String pCollectionA = "pCollectionA";
    final String pTransformIdA = "pTransformIdA";
    MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
    PCollectionConsumerRegistry consumers = new PCollectionConsumerRegistry(metricsContainerRegistry, mock(ExecutionStateTracker.class));
    SplittingReceiver consumerA1 = mock(SplittingReceiver.class);
    consumers.register(pCollectionA, pTransformIdA, consumerA1, StringUtf8Coder.of());
    FnDataReceiver<WindowedValue<String>> wrapperConsumer = (FnDataReceiver<WindowedValue<String>>) (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
    assertTrue(wrapperConsumer instanceof HandlesSplits);
    ((HandlesSplits) wrapperConsumer).getProgress();
    verify(consumerA1).getProgress();
    ((HandlesSplits) wrapperConsumer).trySplit(0.3);
    verify(consumerA1).trySplit(0.3);
}