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_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_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_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_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_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_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_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_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_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_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_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_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(self, mock_process):
action = ActionBalanceRackAware(self.args, self.cluster)
action.process_cluster()
mock_process.assert_has_calls([call(0), call(1)])
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)