@Test
public void getStartingOffsets() throws IOException {

    final Map<Integer, Long> firstOffsets = new HashMap<>();
    final Map<Integer, Long> secondOffsets = new HashMap<>();
    final int firstPartition = 42;
    final int secondPartition = 51;
    final OffsetComputer offsetComputer = mock(OffsetComputer.class);
    final PartitionedWriter<String> writer = new PartitionedWriter<>(mock(BiFunction.class), offsetComputer,
            "ignored", DataAccessEventProtos.FsEvent.newBuilder(), mock(Checkpointer.class));
    firstOffsets.put(firstPartition, 2L);
    firstOffsets.put(secondPartition, 3L);
    secondOffsets.put(firstPartition, 4L);
    secondOffsets.put(secondPartition, 5L);
    when(offsetComputer.computeOffsets(any())).thenReturn(firstOffsets).thenReturn(secondOffsets);
    final List<Integer> bothPartitions = Arrays.asList(firstPartition, secondPartition);
    assertEquals(firstOffsets, writer.getStartingOffsets(bothPartitions));
    assertEquals(firstOffsets, writer.getStartingOffsets(bothPartitions));
    // No unnecessary round-trip if value already filled
    verify(offsetComputer, times(1)).computeOffsets(bothPartitions);
    // If any gets dropped, then everything gets recomputed
    writer.dropPartition(firstPartition);
    assertEquals(secondOffsets, writer.getStartingOffsets(bothPartitions));
}