ZookeeperBrokerStorage implements BrokerStorage { @Override public Multimap<Integer, TopicAndPartition> readLeadersForPartitions(Set<TopicAndPartition> topicAndPartitionSet) { Multimap<Integer, TopicAndPartition> leadersForPartitions = ArrayListMultimap.create(); for (TopicAndPartition topicAndPartition : topicAndPartitionSet) { try { Integer leaderId = readLeaderForPartition(topicAndPartition); leadersForPartitions.put(leaderId, topicAndPartition); } catch (BrokerNotFoundForPartitionException ex) { LOGGER.warn("Broker not found", ex); } } return leadersForPartitions; } @Inject  ZookeeperBrokerStorage(@Named(CuratorType.KAFKA) CuratorFramework curatorFramework,
                                  KafkaZkClient kafkaZkClient, String brokerListenerName); @Override @SuppressWarnings("unchecked") int readLeaderForPartition(TopicAndPartition topicAndPartition); @Override Multimap<Integer, TopicAndPartition> readLeadersForPartitions(Set<TopicAndPartition> topicAndPartitionSet); @Override @SuppressWarnings("unchecked") BrokerDetails readBrokerDetails(Integer brokerId); @Override List<Integer> readPartitionsIds(String topicName);  }
@Test public void shouldReadLeadersForPartitions() throws Exception { TopicAndPartition topicAndPartition1 = new TopicAndPartition("topic1", 2); TopicAndPartition topicAndPartition2 = new TopicAndPartition("topic2", 4); createLeaderForPartition(topicAndPartition1, 0); createLeaderForPartition(topicAndPartition2, 3); Multimap<Integer, TopicAndPartition> leadersForPartitions = brokerStorage.readLeadersForPartitions( ImmutableSet.of(topicAndPartition1, topicAndPartition2) ); assertThat(leadersForPartitions.get(0)).contains(topicAndPartition1); assertThat(leadersForPartitions.get(3)).contains(topicAndPartition2); }
@Test public void shouldNotReadLeadersDueTooNoDataInZk() { Multimap<Integer, TopicAndPartition> leadersForPartitions = brokerStorage.readLeadersForPartitions( ImmutableSet.of(new TopicAndPartition("topic1", 2)) ); assertThat(leadersForPartitions.size()).isZero(); }