def setUp(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
def test_check_partition_swappable_racks_ok(self):
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = Broker(1, 'brokerhost3.example.com')
b4 = Broker(1, 'brokerhost4.example.com')
b2.rack = "a"
b3.rack = "b"
b4.rack = "b"
replicas_a = [b1, b2]
replicas_b = [b3, b4]
assert check_partition_swappable(replicas_a, replicas_b, 0) is True
assert check_partition_swappable(replicas_b, replicas_a, 0) is True
def set_up_cluster_4broker():
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_broker(Broker(2, "brokerhost2.example.com"))
cluster.add_broker(Broker(3, "brokerhost3.example.com"))
cluster.add_broker(Broker(4, "brokerhost4.example.com"))
cluster.brokers[1].rack = "a"
cluster.brokers[2].rack = "a"
cluster.brokers[3].rack = "b"
cluster.brokers[4].rack = "b"
cluster.add_topic(Topic("testTopic1", 4))
cluster.add_topic(Topic("testTopic2", 4))
cluster.add_topic(Topic("testTopic3", 4))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[3], 1)
partition = cluster.topics['testTopic1'].partitions[2]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[3], 1)
partition = cluster.topics['testTopic1'].partitions[3]
partition.add_replica(cluster.brokers[4], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[4], 0)
partition.add_replica(cluster.brokers[3], 1)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[4], 1)
partition = cluster.topics['testTopic2'].partitions[2]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[3]
partition.add_replica(cluster.brokers[3], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic3'].partitions[0]
partition.add_replica(cluster.brokers[3], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic3'].partitions[1]
partition.add_replica(cluster.brokers[4], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic3'].partitions[2]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic3'].partitions[3]
partition.add_replica(cluster.brokers[3], 0)
partition.add_replica(cluster.brokers[4], 1)
return cluster
def test_process_cluster_empty_one(self):
self.cluster.add_broker(Broker(3, 'brokerhost3.example.com'))
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = self.cluster.brokers[3]
self.cluster.add_topic(Topic("testTopic3", 2))
partition = self.cluster.topics['testTopic3'].partitions[0]
partition.add_replica(b3, 0)
partition.add_replica(b2, 1)
partition = self.cluster.topics['testTopic3'].partitions[1]
partition.add_replica(b2, 0)
partition.add_replica(b3, 1)
self.cluster.topics['testTopic1'].partitions[0].swap_replicas(b1, b3)
self.cluster.topics['testTopic1'].partitions[1].swap_replicas(b1, b3)
self.cluster.topics['testTopic2'].partitions[0].swap_replicas(b1, b3)
self.cluster.topics['testTopic2'].partitions[1].swap_replicas(b1, b3)
action = ActionBalanceCount(self.args, self.cluster)
action.process_cluster()
assert len(b1.partitions[0]) == 2
assert len(b1.partitions[1]) == 2
assert len(b2.partitions[0]) == 2
assert len(b2.partitions[1]) == 2
assert len(b3.partitions[0]) == 2
assert len(b3.partitions[1]) == 2
def test_process_cluster_empty_broker(self):
self.cluster.add_broker(Broker(3, 'brokerhost3.example.com'))
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
self.cluster.add_topic(Topic("testTopic3", 2))
partition = self.cluster.topics['testTopic3'].partitions[0]
partition.size = 1000
partition.add_replica(b1, 0)
partition.add_replica(b2, 1)
partition = self.cluster.topics['testTopic3'].partitions[1]
partition.add_replica(b2, 0)
partition.add_replica(b1, 1)
partition.size = 2000
action = ActionBalanceSize(self.args, self.cluster)
action.process_cluster()
assert sum([p.size for p in self.cluster.brokers[1].partitions[0]],
0) == 3000
assert sum([p.size for p in self.cluster.brokers[1].partitions[1]],
0) == 3000
assert sum([p.size for p in self.cluster.brokers[2].partitions[0]],
0) == 3000
assert sum([p.size for p in self.cluster.brokers[2].partitions[1]],
0) == 3000
assert sum([p.size for p in self.cluster.brokers[3].partitions[0]],
0) == 3000
assert sum([p.size for p in self.cluster.brokers[3].partitions[1]],
0) == 3000
def setUp(self):
self.topic = Topic('testTopic', 10)
self.broker = Broker(1, 'brokerhost1.example.com')
for i in range(10):
self.topic.partitions[i].replicas = [self.broker]
self.reassignment = Reassignment(self.topic.partitions, pause_time=0)
self.null_plugin = PluginModule()
def test_partition_copy_with_replicas(self):
broker = Broker(1, 'testhost1')
self.topic.partitions[0].replicas = [broker]
partition2 = self.topic.partitions[0].copy()
assert self.topic.partitions[0] == partition2
assert partition2.replicas == []
assert self.topic.partitions[0] is not partition2
def add_topic_with_replicas(cluster, topic, topic_data):
newtopic = Topic(topic, len(topic_data['partitions']))
for partition in topic_data['partitions']:
for i, replica in enumerate(topic_data['partitions'][partition]):
if replica not in cluster.brokers:
# Hit a replica that's not in the ID list (which means it's dead)
# We'll add it, but trying to get sizes will fail as we don't have a hostname
cluster.add_broker(Broker(replica, None))
newtopic.partitions[int(partition)].add_replica(
cluster.brokers[replica], i)
cluster.add_topic(newtopic)
def test_process_partitions_at_pos_swap_partition(self, mock_pick):
action = ActionBalanceRackAware(self.args, self.cluster)
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = Broker(3, 'brokerhost3.example.com')
b4 = Broker(4, 'brokerhost4.example.com')
self.cluster.add_broker(b3)
self.cluster.add_broker(b4)
b3.rack = "a"
b4.rack = "c"
self.cluster.topics['testTopic2'].partitions[0].swap_replicas(b2, b3)
self.cluster.topics['testTopic1'].partitions[1].swap_replicas(b1, b4)
mock_pick.return_value = self.cluster.topics['testTopic1'].partitions[
1]
action._process_partitions_at_pos(0)
assert self.cluster.topics['testTopic1'].partitions[1].replicas == [
b3, b4
]
assert self.cluster.topics['testTopic2'].partitions[0].replicas == [
b2, b1
]
def test_process_cluster_no_change(self):
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.args.brokers = [3]
self.args.to_broker = 1
action = ActionClone(self.args, self.cluster)
action.process_cluster()
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [b1, b2]
assert self.cluster.topics['testTopic1'].partitions[1].replicas == [b2, b1]
assert self.cluster.topics['testTopic2'].partitions[0].replicas == [b2, b1]
assert self.cluster.topics['testTopic2'].partitions[1].replicas == [b1, b2]
def set_up_cluster():
cluster = Cluster()
cluster.retention = 100000
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_broker(Broker(2, "brokerhost2.example.com"))
cluster.brokers[1].rack = "a"
cluster.brokers[2].rack = "b"
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
return cluster
def create_cluster_onehost(self):
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
return cluster
def test_try_pick_new_broker(self):
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = Broker(3, 'brokerhost3.example.com')
self.cluster.add_broker(b3)
b3.rack = "b"
action = ActionBalanceRackAware(self.args, self.cluster)
# Firmly order the deque
action._random_brokers = deque([b3, b1, b2])
newbroker = action._try_pick_new_broker(
self.cluster.topics['testTopic1'].partitions[0], 1)
assert newbroker == b3
assert action._random_brokers == deque([b1, b2, b3])
def set_up_cluster():
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_broker(Broker(2, "brokerhost2.example.com"))
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition.size = 1001
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition.size = 1002
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition.size = 2001
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition.size = 2002
return cluster
def test_process_cluster(self):
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.args.brokers = [1]
self.args.to_brokers = []
action = ActionRemove(self.args, self.cluster)
action.process_cluster()
b2 = self.cluster.brokers[2]
b3 = self.cluster.brokers[3]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [
b3, b2
]
assert self.cluster.topics['testTopic1'].partitions[1].replicas == [
b2, b3
]
assert self.cluster.topics['testTopic2'].partitions[0].replicas == [
b2, b3
]
assert self.cluster.topics['testTopic2'].partitions[1].replicas == [
b3, b2
]
def test_process_cluster_increase(self):
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.args.topics = ['testTopic1']
self.args.replication_factor = 3
action = ActionSetRF(self.args, self.cluster)
action.process_cluster()
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = self.cluster.brokers[3]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [
b1, b2, b3
]
assert self.cluster.topics['testTopic1'].partitions[1].replicas == [
b2, b1, b3
]
assert self.cluster.topics['testTopic2'].partitions[0].replicas == [
b2, b1
]
assert self.cluster.topics['testTopic2'].partitions[1].replicas == [
b1, b2
]
def test_broker_inequality_id(self):
broker2 = Broker(2, 'brokerhost1.example.com')
assert self.broker != broker2
def test_broker_inequality_hostname(self):
broker2 = Broker(1, 'brokerhost2.example.com')
assert self.broker != broker2
def test_broker_equality(self):
broker2 = Broker(1, 'brokerhost1.example.com')
assert self.broker == broker2
def setUp(self):
self.broker = Broker(1, 'brokerhost1.example.com')
def setUp(self):
self.topic = Topic('testTopic', 10)
self.broker = Broker(1, 'brokerhost1.example.com')
for i in range(10):
self.topic.partitions[i].replicas = [self.broker]
def setUp(self):
self.topic = Topic('testTopic', 10)
self.broker = Broker(1, 'brokerhost1.example.com')
for i in range(10):
self.topic.partitions[i].replicas = [self.broker]
self.replica_election = ReplicaElection(self.topic.partitions, pause_time=0)
def test_partition_equality_with_different_replicas(self):
partition2 = Partition(self.topic, 0)
broker = Broker(1, 'testhost1')
partition2.replicas = [broker]
assert self.topic.partitions[0] == partition2
def test_partition_dict_for_reassignment_with_replicas(self):
broker = Broker(1, 'testhost1')
self.topic.partitions[0].replicas = [broker]
expected = {"topic": 'testTopic', "partition": 0, "replicas": [1]}
assert self.topic.partitions[0].dict_for_reassignment() == expected
def add_brokers(self, num):
for i in range(1, num + 1):
broker = Broker(i, "brokerhost{0}.example.com".format(i))
self.cluster.add_broker(broker)
