@VisibleForTesting
  static List<TopicPartition> getAllTopicPartitions(
      SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> kafkaConsumerFactoryFn,
      Map<String, Object> kafkaConsumerConfig,
      Set<String> topics,
      @Nullable Pattern topicPattern) {
    List<TopicPartition> current = new ArrayList<>();
    try (Consumer<byte[], byte[]> kafkaConsumer =
        kafkaConsumerFactoryFn.apply(kafkaConsumerConfig)) {
      if (topics != null && !topics.isEmpty()) {
        for (String topic : topics) {
          for (PartitionInfo partition : kafkaConsumer.partitionsFor(topic)) {
            current.add(new TopicPartition(topic, partition.partition()));
          }
        }
      } else {
        for (Map.Entry<String, List<PartitionInfo>> topicInfo :
            kafkaConsumer.listTopics().entrySet()) {
          if (topicPattern == null || topicPattern.matcher(topicInfo.getKey()).matches()) {
            for (PartitionInfo partition : topicInfo.getValue()) {
              current.add(new TopicPartition(partition.topic(), partition.partition()));
            }
          }
        }
      }
    }
    return current;
  }
@Test
  public void testGetAllTopicPartitions() throws Exception {
    Consumer<byte[], byte[]> mockConsumer = Mockito.mock(Consumer.class);
    when(mockConsumer.listTopics())
        .thenReturn(
            ImmutableMap.of(
                "topic1",
                ImmutableList.of(
                    new PartitionInfo("topic1", 0, null, null, null),
                    new PartitionInfo("topic1", 1, null, null, null)),
                "topic2",
                ImmutableList.of(
                    new PartitionInfo("topic2", 0, null, null, null),
                    new PartitionInfo("topic2", 1, null, null, null))));
    assertEquals(
        ImmutableList.of(
            new TopicPartition("topic1", 0),
            new TopicPartition("topic1", 1),
            new TopicPartition("topic2", 0),
            new TopicPartition("topic2", 1)),
        WatchForKafkaTopicPartitions.getAllTopicPartitions(
            (input) -> mockConsumer, null, null, null));
  }