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