PartitionedWriter implements Closeable { public void dropPartition(int partitionId) { synchronized (perPartitionDayWriters) { perPartitionDayWriters.remove(partitionId); perPartitionStartOffset.remove(partitionId); } } PartitionedWriter(BiFunction<Integer, LocalDateTime, ExpiringConsumer<MESSAGE_KIND>> writerBuilder,
                             OffsetComputer offsetComputer, String eventName, Message.Builder emptyMessageBuilder,
                             Checkpointer checkpointer); void dropPartition(int partitionId); void write(Instant when, Offset offset, MESSAGE_KIND msg); @Override void close(); Map<Integer, Long> getStartingOffsets(Collection<Integer> partitionsId); void heartbeat(int partition, Offset offset);  }
@Test public void dropPartition() throws IOException { final BiFunction<Integer, LocalDateTime, ExpiringConsumer<String>> writerBuilder = mock(BiFunction.class); final OffsetComputer offsetComputer = new FixedOffsetComputer("0", 0); final PartitionedWriter<String> writer = new PartitionedWriter<>(writerBuilder, offsetComputer, "ignored", DataAccessEventProtos.FsEvent.newBuilder(), mock(Checkpointer.class)); final ExpiringConsumer<String> firstMockConsumer = mock(ExpiringConsumer.class); final ExpiringConsumer<String> secondMockConsumer = mock(ExpiringConsumer.class); final Instant instant = Instant.EPOCH; when(writerBuilder.apply(any(), any(LocalDateTime.class))).thenReturn(firstMockConsumer).thenReturn(secondMockConsumer); writer.write(instant, buildOffset(1, 8), "Ignored"); verify(firstMockConsumer, times(1)).write(anyLong(), anyString(), any(Offset.class)); writer.dropPartition(1); writer.expireConsumers(); writer.write(instant, buildOffset(1, 8), "Ignored"); verify(secondMockConsumer, times(1)).write(anyLong(), anyString(), any(Offset.class)); writer.expireConsumers(); verify(secondMockConsumer, times(1)).isExpired(); verifyNoMoreInteractions(firstMockConsumer); verifyNoMoreInteractions(secondMockConsumer); }