def test_get_partition_info(self):
     partitions_info = rebalance_rackaware.get_partition_info(self.topicInfo_lines, [], self.dummy_topic)
     partition_4 = [element for element in partitions_info if int(element[rebalance_rackaware.PARTITION]) == 4][0]
     self.assertEqual(partition_4[rebalance_rackaware.PARTITION], 4)
     self.assertEqual(partition_4[rebalance_rackaware.LEADER], 1030)
     self.assertEqual(set(partition_4[rebalance_rackaware.REPLICAS]), set([1030, 1001, 1009]))
     self.assertEqual(set(partition_4[rebalance_rackaware.ISR]), set([1030, 1001, 1009]))
 def test_reassignment_plan_HA_for_topic(self):
     host_info = rebalance_rackaware.parse_topo_info(self.cluster_topo, self.brokers_info)
     brokers_replica_count = []
     for host in host_info:
         b = {
             rebalance_rackaware.BROKER_ID : host[rebalance_rackaware.BROKER_ID],
             rebalance_rackaware.LEADERS : 0,
             rebalance_rackaware.FOLLOWERS: 0,
         }
         brokers_replica_count.append(b)
     partitions_info = rebalance_rackaware.get_partition_info(self.topicInfo_lines, [], self.dummy_topic)
     rgen = rebalance_rackaware.ReassignmentGenerator(host_info, self.dummy_topic, partitions_info, self.compute_storage_cost)
     fd_ud_list, fd_list, ud_list = self.generate_fd_ud_list(3,3)
     alternated_list = rgen._generate_alternated_fd_ud_list(fd_ud_list, fd_list, ud_list)
     reassignment_plan, balanced_partitions = rgen._generate_reassignment_plan_for_topic(3,0,alternated_list,3,3,brokers_replica_count, None)
     topic_balanced = rgen._verify_reassignment_plan(reassignment_plan,self.dummy_topic,3,3,3)
     self.assertTrue(topic_balanced)