def test_rf_mismatch(self, create_cluster_topology, assignment, brokers):
with mock.patch.object(ReplaceBrokerCmd,
'process_assignment'), mock.patch.object(
ReplaceBrokerCmd, 'get_topic_filter'):
cmd = ReplaceBrokerCmd()
cmd.args = self.mock_args()
# Remove broker 0
cmd.args.source_broker = 0
cmd.args.dest_broker = None
ct = create_cluster_topology(assignment, brokers)
cb = PartitionCountBalancer(ct, cmd.args)
# No rf mismatch allowed -- should fail
with pytest.raises(SystemExit):
cmd.run_command(ct, cb)
assert cmd.process_assignment.call_count == 0
# Should still fail even if we set rf_change
cmd.args.rf_change = True
ct = create_cluster_topology(assignment, brokers)
with pytest.raises(SystemExit):
cmd.run_command(ct, cb)
assert cmd.process_assignment.call_count == 0
# Allow rf mismatch for T0
ct = create_cluster_topology(assignment, brokers)
cmd.args.rf_mismatch = True
cmd.run_command(ct, cb)
assert cmd.process_assignment.call_count == 1
def test_run_command_remove_replica(self, create_cluster_topology):
assignment = {
(u'T0', 0): ['0', '1', '2', '3', '4'],
(u'T0', 1): ['0', '1', '2', '3', '4'],
}
brokers = {
'0': {
'host': 'host2'
},
'1': {
'host': 'host2'
},
'2': {
'host': 'host3'
},
'3': {
'host': 'host4'
},
'4': {
'host': 'host5'
},
}
with mock.patch.object(SetReplicationFactorCmd, 'process_assignment'):
cmd = SetReplicationFactorCmd()
cmd.args = mock.Mock(spec=Namespace)
cmd.args.topic = u'T0'
cmd.args.replication_factor = 2
cmd.zk = mock.Mock()
cmd.zk.get_topics.side_effect = (lambda topic_id: {
topic_id: {
'partitions': {
'0': {
'isr': ['0', '1']
},
'1': {
'isr': ['1', '2']
},
},
},
})
ct = create_cluster_topology(assignment, brokers)
cb = PartitionCountBalancer(ct, cmd.args)
cmd.run_command(ct, cb)
assert cmd.process_assignment.call_count == 1
args, kwargs = cmd.process_assignment.call_args_list[0]
assignment = args[0]
assert len(assignment) == 2
assert set(assignment[(u'T0', 0)]) == set(['0', '1'])
assert set(assignment[(u'T0', 1)]) == set(['1', '2'])
assert kwargs['allow_rf_change']
def test_filter_list(self, create_cluster_topology, assignment, brokers):
with mock.patch.object(ReplaceBrokerCmd,
'process_assignment'), mock.patch.object(
ReplaceBrokerCmd,
'get_topic_filter') as mock_filter:
cmd = ReplaceBrokerCmd()
cmd.args = self.mock_args()
cmd.args.topic_partition_filter = 'foo'
mock_filter.return_value = {('T0', 0)}
ct = create_cluster_topology(assignment, brokers)
cb = PartitionCountBalancer(ct, cmd.args)
cmd.run_command(ct, cb)
assert cmd.process_assignment.call_count == 1
args, _ = cmd.process_assignment.call_args_list[0]
assignment = args[0]
assert len(assignment) == 1
assert set(assignment[('T0', 0)]) == set([3, 1])
def build_balancer(cluster_topology, **kwargs):
args = mock.Mock(spec=Namespace)
args.balancer_args = []
args.configure_mock(**kwargs)
return PartitionCountBalancer(cluster_topology, args)