Examples with SingularProcessBundleProgressTracker org.apache.beam.runners.dataflow.worker.fn.control.BeamFnMapTaskExecutor.SingularProcessBundleProgressTracker used on opensource projects

Example 1 with SingularProcessBundleProgressTracker

use of org.apache.beam.runners.dataflow.worker.fn.control.BeamFnMapTaskExecutor.SingularProcessBundleProgressTracker in project beam by apache.

the class SingularProcessBundleProgressTrackerTest method testProgressInterpolation.

@Test
public void testProgressInterpolation() throws Exception {
    ReadOperation read = Mockito.mock(ReadOperation.class);
    RemoteGrpcPortWriteOperation grpcWrite = Mockito.mock(RemoteGrpcPortWriteOperation.class);
    RegisterAndProcessBundleOperation process = Mockito.mock(RegisterAndProcessBundleOperation.class);
    when(grpcWrite.processedElementsConsumer()).thenReturn(elementsConsumed -> {
    });
    SingularProcessBundleProgressTracker tracker = new SingularProcessBundleProgressTracker(read, grpcWrite, process);
    when(read.getProgress()).thenReturn(new TestProgress("A"), new TestProgress("B"), new TestProgress("C"));
    when(grpcWrite.getElementsSent()).thenReturn(1, 10, 20, 30);
    // This test ignores them, directly working on mocked getInputElementsConsumed
    when(process.getProcessBundleProgress()).thenReturn(CompletableFuture.completedFuture(BeamFnApi.ProcessBundleProgressResponse.getDefaultInstance()));
    when(process.getInputElementsConsumed(any(Iterable.class))).thenReturn(1L, 4L, 10L).thenThrow(new RuntimeException());
    // Initially no progress is known.
    assertEquals(null, tracker.getWorkerProgress());
    // After reading, and writing, and processing one element, the progress is aligned at A.
    tracker.updateProgress();
    assertEquals(new TestProgress("A"), tracker.getWorkerProgress());
    // We've read up to B (10 elements) but only consumed 4.  Progress remains at A.
    tracker.updateProgress();
    assertEquals(new TestProgress("A"), tracker.getWorkerProgress());
    // Once 10 elements have been consumed, advance to B.
    tracker.updateProgress();
    assertEquals(new TestProgress("B"), tracker.getWorkerProgress());
    // An exception is thrown, default to latest read progress.
    tracker.updateProgress();
    assertEquals(new TestProgress("C"), tracker.getWorkerProgress());
}