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)