def test_generate_alternated_fd_ud_list_test_case_4FDs_6UDs(self):
     (FDs, UDs) = 4,6
     fd_ud_list, fd_list, ud_list = self.generate_fd_ud_list(FDs,UDs)
     rgen = rebalance_rackaware.ReassignmentGenerator("hostInfo", self.dummy_topic, "partitionInfo", self.compute_storage_cost)
     alternated_list = rgen._generate_alternated_fd_ud_list(fd_ud_list, fd_list, ud_list)
     expected_list = ['FD0UD0', 'FD1UD1', 'FD2UD2', 'FD3UD3', 'FD0UD1', 'FD1UD2', 'FD2UD3', 'FD3UD4', 'FD0UD2', 'FD1UD3', 'FD2UD4', 'FD3UD5', 'FD0UD3', 'FD1UD4', 'FD2UD5', 'FD3UD0', 'FD0UD4', 'FD1UD5', 'FD2UD0', 'FD3UD1', 'FD0UD5', 'FD1UD0', 'FD2UD1', 'FD3UD2']
     self.assertEqual(len(alternated_list), len(fd_ud_list))
     self.assertEqual(alternated_list, expected_list)
 def test_generate_alternated_fd_ud_list_test_case_3FDs_3UDs_with_gaps(self):
     fd_ud_list = ['FD0UD1', 'FD0UD2', 'FD1UD0', 'FD1UD1', 'FD2UD0', 'FD2UD2']
     fd_list = ['0', '1', '2']
     ud_list = ['0', '1', '2']
     rgen = rebalance_rackaware.ReassignmentGenerator("hostInfo", self.dummy_topic, "partitionInfo", self.compute_storage_cost)
     alternated_list = rgen._generate_alternated_fd_ud_list(fd_ud_list, fd_list, ud_list)
     expected_list = ['FD1UD1', 'FD2UD2', 'FD0UD1', 'FD2UD0', 'FD0UD2', 'FD1UD0']
     self.assertEqual(len(alternated_list), len(fd_ud_list))
     self.assertEqual(alternated_list, expected_list)
 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)
    def test_check_if_partition_balanced(self):
        host_info = [{'vmId': 1, 'fqdn': 'wn01-foobar', 'updateDomain': '0', 'brokerId': '1016', 'faultDomain': '1', 'rack': 'FD1UD0'},
                    {'vmId': 2, 'fqdn': 'wn02-foobar', 'updateDomain': '2', 'brokerId': '1017', 'faultDomain': '0', 'rack': 'FD0UD2'},
                    {'vmId': 3, 'fqdn': 'wn77-foobar', 'updateDomain': '1', 'brokerId': '1008', 'faultDomain': '2', 'rack': 'FD2UD1'},
                    {'vmId': 4, 'fqdn': 'wn47-foobar', 'updateDomain': '2', 'brokerId': '1098', 'faultDomain': '2', 'rack': 'FD2UD2'},
                    {'vmId': 5, 'fqdn': 'wn56-foobar', 'updateDomain': '1', 'brokerId': '2', 'faultDomain': '2', 'rack': 'FD2UD1'}]
        rgen = rebalance_rackaware.ReassignmentGenerator(host_info, self.dummy_topic, "partitions", self.compute_storage_cost)

        partition_1 = {"partition": 1,"leader": 1016,"replicas": [1016, 1017, 1008],"isr": [1016, 1017, 1008]}
        partition_2 = {"partition": 1,"leader": 2,"replicas": [1016, 2, 1098],"isr": [1016]}
        partition_3 = {"partition": 1,"leader": 2,"replicas": [1016, 2],"isr": [1016]}
        partition_4 = {"partition": 1,"leader": 2,"replicas": [1016, 1017, 1008, 1098],"isr": [1016]}
        partition_5 = {"partition": 1,"leader": 2,"replicas": [1008, 1098],"isr": [1008]}
        partition_6 = {"partition": 1,"leader": 2,"replicas": [1016, 1017],"isr": [1008]}
        
        #Balanced Checks
        self.assertTrue(rgen._check_if_partition_balanced(partition_1,3,3,3,[],[]))
        self.assertTrue(rgen._check_if_partition_balanced(partition_3,2,3,3,[],[]))
        self.assertTrue(rgen._check_if_partition_balanced(partition_4,4,3,3,[],[]))

        #Imbalanced Checks
        self.assertFalse(rgen._check_if_partition_balanced(partition_2,3,3,3,[],[]))
        self.assertFalse(rgen._check_if_partition_balanced(partition_5,2,3,3,[],[]))
        self.assertFalse(rgen._check_if_partition_balanced(partition_6,3,3,3,[],[])) 
 def test_generate_fd_ud_list(self):
     host_info = [{'vmId': 37, 'fqdn': 'wn25-foobar', 'updateDomain': '1', 'brokerId': '1016', 'faultDomain': '1', 'rack': 'FD1UD1'},{'vmId': 1, 'fqdn': 'wn25-foobar', 'updateDomain': '1', 'brokerId': '1016', 'faultDomain': '1', 'rack': 'FD1UD1'},{'vmId': 10, 'fqdn': 'wn7-foobar', 'updateDomain': '1', 'brokerId': '1008', 'faultDomain': '2', 'rack': 'FD2UD1'}]   
     rgen = rebalance_rackaware.ReassignmentGenerator(host_info, self.dummy_topic, "partitionInfo", self.compute_storage_cost)
     fd_ud_list = rgen._generate_fd_ud_list()
     self.assertEqual(fd_ud_list, ['FD1UD1', 'FD2UD1'])