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)