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

Example 6 with ProcessContinuation

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

the class WatchTest method testNonPollingGrowthTrackerIgnoresWatermark.

@Test
public void testNonPollingGrowthTrackerIgnoresWatermark() throws Exception {
    Instant now = Instant.now();
    PollResult<String> claim = PollResult.incomplete(Arrays.asList(TimestampedValue.of("d", now.plus(standardSeconds(4))), TimestampedValue.of("c", now.plus(standardSeconds(3))), TimestampedValue.of("a", now.plus(standardSeconds(1))), TimestampedValue.of("b", now.plus(standardSeconds(2))))).withWatermark(now.plus(standardSeconds(7)));
    Watch.Growth<String, String, String> growth = Watch.growthOf(new Watch.Growth.PollFn<String, String>() {

        @Override
        public PollResult<String> apply(String element, Context c) throws Exception {
            fail("Never expected to be invoked for NonPollingGrowthState.");
            return null;
        }
    }).withPollInterval(standardSeconds(10));
    GrowthTracker<String, Integer> tracker = newTracker(NonPollingGrowthState.of(claim));
    WatchGrowthFn<String, String, String, Integer> growthFn = new WatchGrowthFn(growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
    DoFn.ProcessContext context = mock(DoFn.ProcessContext.class);
    ManualWatermarkEstimator<Instant> watermarkEstimator = new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
    ProcessContinuation processContinuation = growthFn.process(context, tracker, watermarkEstimator);
    assertEquals(BoundedWindow.TIMESTAMP_MIN_VALUE, watermarkEstimator.currentWatermark());
    assertFalse(processContinuation.shouldResume());
}

Example 7 with ProcessContinuation

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

the class ChildPartitionsRecordActionTest method testRestrictionClaimedAnsIsSplitCaseAndChildExists.

@Test
public void testRestrictionClaimedAnsIsSplitCaseAndChildExists() {
    final String partitionToken = "partitionToken";
    final long heartbeat = 30L;
    final Timestamp startTimestamp = Timestamp.ofTimeMicroseconds(10L);
    final Timestamp endTimestamp = Timestamp.ofTimeMicroseconds(20L);
    final PartitionMetadata partition = mock(PartitionMetadata.class);
    final ChildPartitionsRecord record = new ChildPartitionsRecord(startTimestamp, "recordSequence", Arrays.asList(new ChildPartition("childPartition1", partitionToken), new ChildPartition("childPartition2", partitionToken)), null);
    when(partition.getEndTimestamp()).thenReturn(endTimestamp);
    when(partition.getHeartbeatMillis()).thenReturn(heartbeat);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    when(tracker.tryClaim(10L)).thenReturn(true);
    when(transaction.getPartition("childPartition1")).thenReturn(mock(Struct.class));
    when(transaction.getPartition("childPartition2")).thenReturn(mock(Struct.class));
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, watermarkEstimator);
    assertEquals(Optional.empty(), maybeContinuation);
    verify(watermarkEstimator).setWatermark(new Instant(startTimestamp.toSqlTimestamp().getTime()));
}

Example 8 with ProcessContinuation

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

the class ChildPartitionsRecordActionTest method testRestrictionNotClaimed.

@Test
public void testRestrictionNotClaimed() {
    final String partitionToken = "partitionToken";
    final Timestamp startTimestamp = Timestamp.ofTimeMicroseconds(10L);
    final PartitionMetadata partition = mock(PartitionMetadata.class);
    final ChildPartitionsRecord record = new ChildPartitionsRecord(startTimestamp, "recordSequence", Arrays.asList(new ChildPartition("childPartition1", partitionToken), new ChildPartition("childPartition2", partitionToken)), null);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    when(tracker.tryClaim(10L)).thenReturn(false);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, watermarkEstimator);
    assertEquals(Optional.of(ProcessContinuation.stop()), maybeContinuation);
    verify(watermarkEstimator, never()).setWatermark(any());
    verify(dao, never()).insert(any());
}

Example 9 with ProcessContinuation

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

the class HeartbeatRecordActionTest method testRestrictionClaimed.

@Test
public void testRestrictionClaimed() {
    final String partitionToken = "partitionToken";
    final Timestamp timestamp = Timestamp.ofTimeMicroseconds(10L);
    when(tracker.tryClaim(10L)).thenReturn(true);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, new HeartbeatRecord(timestamp, null), tracker, watermarkEstimator);
    assertEquals(Optional.empty(), maybeContinuation);
    verify(watermarkEstimator).setWatermark(new Instant(timestamp.toSqlTimestamp().getTime()));
}

Example 10 with ProcessContinuation

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

the class HeartbeatRecordActionTest method testRestrictionNotClaimed.

@Test
public void testRestrictionNotClaimed() {
    final String partitionToken = "partitionToken";
    final Timestamp timestamp = Timestamp.ofTimeMicroseconds(10L);
    when(tracker.tryClaim(10L)).thenReturn(false);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, new HeartbeatRecord(timestamp, null), tracker, watermarkEstimator);
    assertEquals(Optional.of(ProcessContinuation.stop()), maybeContinuation);
    verify(watermarkEstimator, never()).setWatermark(any());
}