protected Map<String, Integer> getNumPartitions(final Set<String> topics,
                                                    final Set<String> tempUnknownTopics) {
        log.debug("Trying to check if topics {} have been created with expected number of partitions.", topics);

        final Map<String, List<TopicPartitionInfo>> topicPartitionInfo = getTopicPartitionInfo(topics, tempUnknownTopics);
        return topicPartitionInfo.entrySet().stream().collect(Collectors.toMap(
            Entry::getKey, e -> e.getValue().size()));
    }
@Test
    public void shouldReturnCorrectPartitionCounts() {
        mockAdminClient.addTopic(
            false,
            topic1,
            Collections.singletonList(new TopicPartitionInfo(0, broker1, singleReplica, Collections.emptyList())),
            null);
        assertEquals(Collections.singletonMap(topic1, 1),
            internalTopicManager.getNumPartitions(Collections.singleton(topic1), Collections.emptySet()));
    }