def test_run_command_add_replica(self, create_cluster_topology):
assignment = {
(u'T0', 0): ['0', '1'],
(u'T0', 1): ['1', '2'],
}
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 = 5
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', '2', '3', '4'])
assert set(assignment[(u'T0', 1)]) == set(['0', '1', '2', '3', '4'])
assert kwargs['allow_rf_change']
def test_run_command_add_replica(self, create_cluster_topology):
assignment = {
(u'T0', 0): ['0', '1'],
(u'T0', 1): ['1', '2'],
}
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 = 5
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', '2', '3', '4'])
assert set(assignment[(u'T0',
1)]) == set(['0', '1', '2', '3', '4'])
assert kwargs['allow_rf_change']
def parse_args():
"""Parse the arguments."""
parser = argparse.ArgumentParser(
description='Manage and describe partition layout over brokers of'
' a cluster.', )
parser.add_argument(
'--cluster-type',
'-t',
dest='cluster_type',
help='Type of the cluster.',
type=str,
required=True,
)
parser.add_argument(
'--cluster-name',
'-c',
dest='cluster_name',
help='Name of the cluster (Default to local cluster).',
)
parser.add_argument(
'--discovery-base-path',
dest='discovery_base_path',
type=str,
help='Path of the directory containing the <cluster_type>.yaml config',
)
parser.add_argument(
'--logconf',
type=str,
help='Path to logging configuration file. Default: log to console.',
)
parser.add_argument(
'--apply',
action='store_true',
help='Proposed-plan will be executed on confirmation.',
)
parser.add_argument(
'--no-confirm',
action='store_true',
help='Proposed-plan will be executed without confirmation.'
' --apply flag also required.',
)
parser.add_argument(
'--write-to-file',
dest='proposed_plan_file',
metavar='<reassignment-plan-file-path>',
type=str,
help='Write the partition reassignment plan '
'to a json file.',
)
parser.add_argument(
'--group-parser',
type=str,
help='Module containing an implementation of ReplicationGroupParser. '
'The module should be specified as path_to_include_to_py_path:module. '
'Ex: "/module/path:module.parser". '
'If not specified the default replication group parser will create '
'only one group for all brokers.',
)
parser.add_argument(
'--partition-measurer',
type=str,
help='Module containing an implementation of PartitionMeasurer. '
'The module should be specified as path_to_include_to_py_path:module. '
'Default: Assign each partition a weight and size of 1.')
parser.add_argument(
'--measurer-args',
type=str,
action='append',
default=[],
help='Argument list that is passed to the chosen PartitionMeasurer. '
'Ex: --measurer-args "--n 10" will pass ["--n", "10"] to the '
'PartitionMeasurer\'s parse_args method.')
parser.add_argument(
'--cluster-balancer',
type=str,
help='Module containing an implementation of ClusterBalancer. '
'The module should be specified as path_to_include_to_py_path:module. '
'Default: PartitionCountBalancer.',
)
parser.add_argument(
'--balancer-args',
type=str,
action='append',
default=[],
help='Argument list that is passed to the chosen ClusterBalancer. '
'Ex: --balancer-args "--n 10" will pass ["--n", "10"] to the '
'ClusterBalancer\'s parse_args method.')
parser.add_argument(
'--partition-count-balancer',
action='store_const',
const=PARTITION_COUNT_BALANCER_MODULE,
dest='cluster_balancer',
help='Use the number of partitions on each broker to balance the '
'cluster.',
)
parser.add_argument(
'--genetic-balancer',
action='store_const',
const=GENETIC_BALANCER_MODULE,
dest='cluster_balancer',
help='Use partition metrics and a genetic algorithm to balance the '
'cluster.',
)
subparsers = parser.add_subparsers()
RebalanceCmd().add_subparser(subparsers)
DecommissionCmd().add_subparser(subparsers)
RevokeLeadershipCmd().add_subparser(subparsers)
StatsCmd().add_subparser(subparsers)
StoreAssignmentsCmd().add_subparser(subparsers)
ReplaceBrokerCmd().add_subparser(subparsers)
SetReplicationFactorCmd().add_subparser(subparsers)
PreferredReplicaElectionCmd().add_subparser(subparsers)
return parser.parse_args()