for (boolean hasNext = reader.start(); hasNext; ) { actual.add(reader.getCurrent()); Checkpoint<T> checkpoint = reader.getCheckpointMark(); Coder<Checkpoint<T>> checkpointCoder = unboundedSource.getCheckpointMarkCoder(); Checkpoint<T> decodedCheckpoint = CoderUtils.decodeFromByteArray( checkpointCoder, CoderUtils.encodeToByteArray(checkpointCoder, checkpoint)); reader.close(); checkpoint.finalizeCheckpoint(); unboundedSource.createReader(options, decodedCheckpoint); reader = restarted; hasNext = reader.start(); } else { hasNext = reader.advance(); Checkpoint<T> checkpointDone = reader.getCheckpointMark(); assertTrue( checkpointDone.getResidualElements() == null
private <T> void testBoundedToUnboundedSourceAdapterCheckpoint( BoundedSource<T> boundedSource, List<T> expectedElements) throws Exception { BoundedToUnboundedSourceAdapter<T> unboundedSource = new BoundedToUnboundedSourceAdapter<>(boundedSource); PipelineOptions options = PipelineOptionsFactory.create(); BoundedToUnboundedSourceAdapter<T>.Reader reader = unboundedSource.createReader(options, null); List<T> actual = Lists.newArrayList(); for (boolean hasNext = reader.start(); hasNext; hasNext = reader.advance()) { actual.add(reader.getCurrent()); // checkpoint every 9 elements if (actual.size() % 9 == 0) { Checkpoint<T> checkpoint = reader.getCheckpointMark(); checkpoint.finalizeCheckpoint(); } } Checkpoint<T> checkpointDone = reader.getCheckpointMark(); assertTrue( checkpointDone.getResidualElements() == null || checkpointDone.getResidualElements().isEmpty()); assertEquals(expectedElements.size(), actual.size()); assertEquals(Sets.newHashSet(expectedElements), Sets.newHashSet(actual)); }
@Test public void testReadFromCheckpointBeforeStart() throws Exception { thrown.expect(NoSuchElementException.class); BoundedSource<Long> countingSource = CountingSource.upTo(100); BoundedToUnboundedSourceAdapter<Long> unboundedSource = new BoundedToUnboundedSourceAdapter<>(countingSource); PipelineOptions options = PipelineOptionsFactory.create(); List<TimestampedValue<Long>> elements = ImmutableList.of(TimestampedValue.of(1L, new Instant(1L))); Checkpoint<Long> checkpoint = new Checkpoint<>(elements, countingSource); unboundedSource.createReader(options, checkpoint).getCurrent(); }
@Test public void testReadBeforeStart() throws Exception { thrown.expect(NoSuchElementException.class); BoundedSource<Long> countingSource = CountingSource.upTo(100); BoundedToUnboundedSourceAdapter<Long> unboundedSource = new BoundedToUnboundedSourceAdapter<>(countingSource); PipelineOptions options = PipelineOptionsFactory.create(); unboundedSource.createReader(options, null).getCurrent(); }
@Override public Reader createReader(PipelineOptions options, Checkpoint<T> checkpoint) throws IOException { if (checkpoint == null) { return new Reader(null /* residualElements */, boundedSource, options); } else { return new Reader(checkpoint.residualElements, checkpoint.residualSource, options); } }