@Override
  protected Source.Reader<T> createReader(@Nonnull FlinkSourceSplit<T> sourceSplit)
      throws IOException {
    Source<T> beamSource = sourceSplit.getBeamSplitSource();
    return ((BoundedSource<T>) beamSource).createReader(pipelineOptions);
  }
@Test
  public void testSnapshotStateAndRestore() throws Exception {
    final int numSplits = 2;
    final int numRecordsPerSplit = 10;

    List<FlinkSourceSplit<KV<Integer, Integer>>> splits =
        createSplits(numSplits, numRecordsPerSplit, 0);
    RecordsValidatingOutput validatingOutput = new RecordsValidatingOutput(splits);
    List<FlinkSourceSplit<KV<Integer, Integer>>> snapshot;

    // Create a reader, take a snapshot.
    try (SourceReader<WindowedValue<KV<Integer, Integer>>, FlinkSourceSplit<KV<Integer, Integer>>>
        reader = createReader()) {
      // Only poll half of the records in the first split.
      pollAndValidate(reader, splits, validatingOutput, numRecordsPerSplit / 2);
      snapshot = reader.snapshotState(0L);
    }

    // Create another reader, add the snapshot splits back.
    try (SourceReader<WindowedValue<KV<Integer, Integer>>, FlinkSourceSplit<KV<Integer, Integer>>>
        reader = createReader()) {
      pollAndValidate(reader, snapshot, validatingOutput, Integer.MAX_VALUE);
    }
  }