HeartbeatConsumer implements GarmadonReader.GarmadonMessageHandler { public void dropPartition(int partition) { latestPartitionsOffset.remove(partition); latestHeartbeats.remove(partition); } HeartbeatConsumer(Collection<PartitionedWriter<MESSAGE_KIND>> writers, TemporalAmount period); void run(); @Override void handle(GarmadonMessage msg); void dropPartition(int partition);  }
@Test public void dropPartition() throws InterruptedException { final Offset firstPartitionOffset = new TopicPartitionOffset(TOPIC, BASE_PARTITION, BASE_OFFSET + 1); final Offset secondPartitionOffset = new TopicPartitionOffset(TOPIC, BASE_PARTITION + 1, BASE_OFFSET + 11); final PartitionedWriter<String> writer = mock(PartitionedWriter.class); final HeartbeatConsumer hb = new HeartbeatConsumer<>(Collections.singleton(writer), Duration.ofSeconds(2)); hb.handle(buildGarmadonMessage(firstPartitionOffset)); hb.handle(buildGarmadonMessage(secondPartitionOffset)); hb.dropPartition(BASE_PARTITION); hb.run(); Thread.sleep(1000); verify(writer, times(1)).heartbeat(eq(BASE_PARTITION + 1), argThat(new OffsetArgumentMatcher(secondPartitionOffset))); verify(writer, never()).heartbeat(eq(BASE_PARTITION), any(Offset.class)); }