Examples with ProcessContinuation org.apache.beam.sdk.transforms.DoFn.ProcessContinuation used on opensource projects

Example 21 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class ReadChangeStreamPartitionDoFnTest method testQueryChangeStreamMode.

@Test
public void testQueryChangeStreamMode() {
    when(queryChangeStreamAction.run(any(), any(), any(), any(), any())).thenReturn(ProcessContinuation.stop());
    final ProcessContinuation result = doFn.processElement(partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
    assertEquals(ProcessContinuation.stop(), result);
    verify(queryChangeStreamAction).run(partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
    verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any());
    verify(heartbeatRecordAction, never()).run(any(), any(), any(), any());
    verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any());
    verify(restrictionTracker, never()).tryClaim(any());
}

Example 22 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class ReadFromKafkaDoFnTest method testProcessElement.

@Test
public void testProcessElement() throws Exception {
    MockOutputReceiver receiver = new MockOutputReceiver();
    consumer.setNumOfRecordsPerPoll(3L);
    long startOffset = 5L;
    OffsetRangeTracker tracker = new OffsetRangeTracker(new OffsetRange(startOffset, startOffset + 3));
    KafkaSourceDescriptor descriptor = KafkaSourceDescriptor.of(topicPartition, null, null, null, null, null);
    ProcessContinuation result = dofnInstance.processElement(descriptor, tracker, null, (OutputReceiver) receiver);
    assertEquals(ProcessContinuation.stop(), result);
    assertEquals(createExpectedRecords(descriptor, startOffset, 3, "key", "value"), receiver.getOutputs());
}