@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 testGetAllTopicPartitionsWithGivenPattern() 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)),
                "topicA",
                ImmutableList.of(
                    new PartitionInfo("topicA", 0, null, null, null),
                    new PartitionInfo("topicA", 1, null, null, null)),
                "topicB",
                ImmutableList.of(
                    new PartitionInfo("topicB", 0, null, null, null),
                    new PartitionInfo("topicB", 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, Pattern.compile("topic[0-9]")));
    assertEquals(
        ImmutableList.of(
            new TopicPartition("topicA", 0),
            new TopicPartition("topicA", 1),
            new TopicPartition("topicB", 0),
            new TopicPartition("topicB", 1)),
        WatchForKafkaTopicPartitions.getAllTopicPartitions(
            (input) -> mockConsumer, null, null, Pattern.compile("topic[A-Z]")));
  }