Examples with OffsetRange org.apache.beam.sdk.io.range.OffsetRange used on opensource projects

Example 31 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange 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());
}

Example 32 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class ReadFromKafkaDoFnTest method testInitialRestrictionWithConsumerPosition.

@Test
public void testInitialRestrictionWithConsumerPosition() throws Exception {
    long expectedStartOffset = 5L;
    consumer.setCurrentPos(5L);
    OffsetRange result = dofnInstance.initialRestriction(KafkaSourceDescriptor.of(topicPartition, null, null, null, null, ImmutableList.of()));
    assertEquals(new OffsetRange(expectedStartOffset, Long.MAX_VALUE), result);
}

Example 33 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class ReadFromKafkaDoFnTest method testInitialRestrictionWhenHasStartTime.

@Test
public void testInitialRestrictionWhenHasStartTime() throws Exception {
    long expectedStartOffset = 10L;
    Instant startReadTime = Instant.now();
    consumer.setStartOffsetForTime(expectedStartOffset, startReadTime);
    consumer.setCurrentPos(5L);
    OffsetRange result = dofnInstance.initialRestriction(KafkaSourceDescriptor.of(topicPartition, null, startReadTime, null, null, ImmutableList.of()));
    assertEquals(new OffsetRange(expectedStartOffset, Long.MAX_VALUE), result);
}

Example 34 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class ReadFromKafkaDoFnTest method testInitialRestrictionWhenHasStopTime.

@Test
public void testInitialRestrictionWhenHasStopTime() throws Exception {
    long expectedStartOffset = 10L;
    Instant startReadTime = Instant.now();
    long expectedStopOffset = 100L;
    Instant stopReadTime = startReadTime.plus(org.joda.time.Duration.millis(2000));
    consumer.setStartOffsetForTime(expectedStartOffset, startReadTime);
    consumer.setStopOffsetForTime(expectedStopOffset, stopReadTime);
    consumer.setCurrentPos(5L);
    OffsetRange result = dofnInstance.initialRestriction(KafkaSourceDescriptor.of(topicPartition, null, startReadTime, null, stopReadTime, ImmutableList.of()));
    assertEquals(new OffsetRange(expectedStartOffset, expectedStopOffset), result);
}

Example 35 with OffsetRange

use of org.apache.beam.sdk.io.range.OffsetRange in project beam by apache.

the class ParquetIOTest method testBlockTracker.

@Test
public void testBlockTracker() {
    OffsetRange range = new OffsetRange(0, 1);
    ParquetIO.ReadFiles.BlockTracker tracker = new ParquetIO.ReadFiles.BlockTracker(range, 7, 3);
    assertEquals(tracker.getProgress().getWorkRemaining(), 1.0, 0.01);
    assertEquals(tracker.getProgress().getWorkCompleted(), 0.0, 0.01);
    tracker.tryClaim(0L);
    tracker.tryClaim(1L);
    assertEquals(tracker.getProgress().getWorkRemaining(), 0.0, 0.01);
    assertEquals(tracker.getProgress().getWorkCompleted(), 1.0, 0.01);
}