def test_try_pick_swap_partition_full(self, mock_check): action = ActionBalanceRackAware(self.args, self.cluster) small_partitions = [ self.cluster.topics['testTopic1'].partitions[0], self.cluster.topics['testTopic1'].partitions[1] ] large_partitions = [self.cluster.topics['testTopic2'].partitions[0]] p = self.cluster.topics['testTopic2'].partitions[1] small_partitions[0].size = 1000 small_partitions[1].size = 2000 large_partitions[0].size = 4000 p.size = 3000 mock_check.side_effect = [False, False, True] target = action._try_pick_swap_partition(p, 0, small_partitions, large_partitions) b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] mock_check.assert_has_calls([ call([b1, b2], [b2, b1], 0), call([b1, b2], [b2, b1], 0), call([b1, b2], [b1, b2], 0) ]) assert target == self.cluster.topics['testTopic1'].partitions[0]
def test_get_sorted_list(self): action = ActionBalanceRackAware(self.args, self.cluster) p1 = self.cluster.topics['testTopic1'].partitions[0] p2 = self.cluster.topics['testTopic1'].partitions[1] p3 = self.cluster.topics['testTopic2'].partitions[0] p4 = self.cluster.topics['testTopic2'].partitions[1] p1.size = 1000 p2.size = 6000 p3.size = 3000 p4.size = 4000 assert action._get_sorted_partition_list_at_pos(0) == [p1, p3, p4, p2]
def test_try_pick_swap_partitions_nolarge(self, mock_check): action = ActionBalanceRackAware(self.args, self.cluster) large_partitions = [] small_partitions = [self.cluster.topics['testTopic2'].partitions[0]] p = self.cluster.topics['testTopic2'].partitions[1] small_partitions[0].size = 1000 p.size = 3000 mock_check.return_value = True assert action._try_pick_swap_partition( p, 0, small_partitions, large_partitions ) == self.cluster.topics['testTopic2'].partitions[0]
def test_try_pick_swap_partition_none(self): action = ActionBalanceRackAware(self.args, self.cluster) small_partitions = [ self.cluster.topics['testTopic1'].partitions[0], self.cluster.topics['testTopic1'].partitions[1] ] large_partitions = [self.cluster.topics['testTopic2'].partitions[0]] p = self.cluster.topics['testTopic2'].partitions[1] small_partitions[0].size = 1000 small_partitions[1].size = 2000 large_partitions[0].size = 4000 p.size = 3000 assert action._try_pick_swap_partition(p, 0, small_partitions, large_partitions) is None
def test_try_pick_new_broker(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(3, 'brokerhost3.example.com') self.cluster.add_broker(b3) b3.rack = "b" action = ActionBalanceRackAware(self.args, self.cluster) # Firmly order the deque action._random_brokers = deque([b3, b1, b2]) newbroker = action._try_pick_new_broker( self.cluster.topics['testTopic1'].partitions[0], 1) assert newbroker == b3 assert action._random_brokers == deque([b1, b2, b3])
def test_process_partitions_at_pos_swap_partition(self, mock_pick): action = ActionBalanceRackAware(self.args, self.cluster) b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(3, 'brokerhost3.example.com') b4 = Broker(4, 'brokerhost4.example.com') self.cluster.add_broker(b3) self.cluster.add_broker(b4) b3.rack = "a" b4.rack = "c" self.cluster.topics['testTopic2'].partitions[0].swap_replicas(b2, b3) self.cluster.topics['testTopic1'].partitions[1].swap_replicas(b1, b4) mock_pick.return_value = self.cluster.topics['testTopic1'].partitions[ 1] action._process_partitions_at_pos(0) assert self.cluster.topics['testTopic1'].partitions[1].replicas == [ b3, b4 ] assert self.cluster.topics['testTopic2'].partitions[0].replicas == [ b2, b1 ]
def test_init_broker_deque(self): action = ActionBalanceRackAware(self.args, self.cluster) assert set(action._random_brokers) == set( self.cluster.brokers.values())
def test_process_cluster(self, mock_process): action = ActionBalanceRackAware(self.args, self.cluster) action.process_cluster() mock_process.assert_has_calls([call(0), call(1)])
def test_process_cluster_single_rack(self): action = ActionBalanceRackAware(self.args, self.cluster) self.cluster.brokers[2].rack = "a" self.assertRaises(BalanceException, action.process_cluster)
def test_process_partitions_at_pos_nochange(self, mock_pick): action = ActionBalanceRackAware(self.args, self.cluster) action._process_partitions_at_pos(0) mock_pick.assert_not_called()
def test_try_pick_new_broker_failed(self): action = ActionBalanceRackAware(self.args, self.cluster) self.assertRaises(BalanceException, action._try_pick_new_broker, self.cluster.topics['testTopic1'].partitions[0], 0)