def test_rebalance_empty_brokers_and_exclude_topics(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['2', '3'],
('t1', '0'): ['2', '3'],
('t1', '1'): ['3', '4'],
('t1', '2'): ['4', '5'],
('t2', '0'): ['3', '4'],
('t2', '1'): ['4', '5'],
('t2', '2'): ['5', '6'],
}
brokers, zk = self._create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers, ['2', '3'], ['t1'])
while o.run([]):
pass
assert distribution[('t1', '0')] == ['2', '3']
assert distribution[('t1', '1')] == ['3', '4']
assert distribution[('t1', '2')] == ['4', '5']
distribution.pop(('t1', '0'))
distribution.pop(('t1', '1'))
distribution.pop(('t1', '2'))
brokers = [item for sublist in distribution.values() for item in sublist]
assert '2' not in brokers
assert '3' not in brokers
def test_rebalance_empty_brokers_and_exclude_topics(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['2', '3'],
('t1', '0'): ['2', '3'],
('t1', '1'): ['3', '4'],
('t1', '2'): ['4', '5'],
('t2', '0'): ['3', '4'],
('t2', '1'): ['4', '5'],
('t2', '2'): ['5', '6'],
}
brokers, zk = self._create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers, ['2', '3'], ['t1'])
while o.run([]):
pass
assert distribution[('t1', '0')] == ['2', '3']
assert distribution[('t1', '1')] == ['3', '4']
assert distribution[('t1', '2')] == ['4', '5']
distribution.pop(('t1', '0'))
distribution.pop(('t1', '1'))
distribution.pop(('t1', '2'))
brokers = [
item for sublist in distribution.values() for item in sublist
]
assert '2' not in brokers
assert '3' not in brokers
def test_rebalance_recovered_with_additional_copy2(self):
distribution = {
('t0', '0'): ['2', '1'],
('t0', '1'): ['1', '2'],
('t0', '2'): ['3', '4']
}
_, zk = _create_zk_for_topics(distribution, ['1', '2', '4'])
o = OptimizedRebalanceChange(zk, ['1', '2', '4'])
while o.run([]):
pass
_verify_balanced(['1', '2', '4'], distribution)
def test_leader_partition_limit(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['1', '2'],
('t0', '2'): ['1', '2'],
('t1', '2'): ['1', '2'],
}
_, zk = _create_zk_for_topics(distribution, ['2', '3'])
o = OptimizedRebalanceChange(zk, ['2', '3'])
while o.run([]):
pass
_verify_balanced(['2', '3'], distribution)
def test_rebalance_empty_one_broker(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['2', '3'],
('t1', '0'): ['2', '3'],
('t1', '1'): ['3', '4'],
}
brokers, zk = _create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers, ['2'], [])
while o.run([]):
pass
_verify_empty_brokers(('2'), distribution)
def test_rebalance_on_filled1(self):
distribution = {
('t0', '0'): ['2'],
('t0', '1'): ['1'],
('t0', '2'): ['1'],
('t0', '3'): ['1'],
}
brokers, zk = _create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers)
# broker to partitions
while o.run([]):
pass
_verify_balanced(('1', '2'), distribution)
def test_rebalance_with_many_topics(self):
distribution = {}
topic_count = 1000
partition_count = 8
for i in range(0, topic_count):
topic = 't{}'.format(i)
distribution.update({(topic, str(partition)): ['1', '2', '3'] for partition in range(0, partition_count)})
_, zk = _create_zk_for_topics(distribution, broker_ids=['1', '2', '3', '4', '5'])
o = OptimizedRebalanceChange(zk, ['1', '2', '3', '4', '5'])
steps = 0
while o.run([]):
steps += 1
_verify_balanced(['1', '2', '3', '4', '5'], distribution)
def test_rebalance_fail_with_not_enough_replicas(self):
distribution = {
('t0', '0'): ['2', '1', '3'],
('t0', '1'): ['1', '2'],
}
_, zk = _create_zk_for_topics(distribution, broker_ids=['1', '3'])
o = OptimizedRebalanceChange(zk, ['1', '3'])
try:
while o.run([]):
pass
assert False, "Balancing can not work with low replication factor"
except Exception:
pass
def test_rebalance_with_dead_brokers(self):
distribution = {
('t0', '0'): ['2', '1'],
('t0', '1'): ['1', '2'],
('t0', '2'): ['1', '2'],
('t0', '3'): ['1', '2'],
('t0', '4'): ['1', '2'],
('t0', '5'): ['1', '2'],
('t0', '6'): ['1', '2'],
}
_, zk = _create_zk_for_topics(distribution, broker_ids=['1', '3'])
o = OptimizedRebalanceChange(zk, ['1', '3'])
while o.run([]):
pass
_verify_balanced(['1', '3'], distribution)
def test_rebalance_on_empty1(self):
brokers, zk = _create_zk_for_topics({})
o = OptimizedRebalanceChange(zk, brokers)
while o.run([]):
pass