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 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 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_process_cluster_large_partition_early(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] self.cluster.add_topic(Topic("testTopic3", 3)) 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(b1, 0) partition.add_replica(b2, 1) partition.size = 2000 partition = self.cluster.topics['testTopic3'].partitions[2] partition.add_replica(b1, 0) partition.add_replica(b2, 1) partition.size = 1000 self.cluster.topics['testTopic1'].partitions[0].size = 8000 action = ActionBalanceSize(self.args, self.cluster) action.process_cluster() b1_0 = sum([p.size for p in self.cluster.brokers[1].partitions[0]], 0) b1_1 = sum([p.size for p in self.cluster.brokers[1].partitions[1]], 0) b2_0 = sum([p.size for p in self.cluster.brokers[2].partitions[0]], 0) b2_1 = sum([p.size for p in self.cluster.brokers[2].partitions[1]], 0) assert b1_0 >= 8000 and b1_0 <= 9000 assert b1_1 >= 8000 and b1_1 <= 9000 assert b2_0 >= 8000 and b2_0 <= 9000 assert b2_1 >= 8000 and b2_1 <= 9000
def test_process_cluster_odd_partitions(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] self.cluster.add_topic(Topic("testTopic3", 3)) 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 partition = self.cluster.topics['testTopic3'].partitions[2] partition.add_replica(b2, 0) partition.add_replica(b1, 1) partition.size = 1000 action = ActionBalanceSize(self.args, self.cluster) action.process_cluster() assert sum([p.size for p in self.cluster.brokers[1].partitions[0]], 0) == 5000 assert sum([p.size for p in self.cluster.brokers[1].partitions[1]], 0) == 5000 assert sum([p.size for p in self.cluster.brokers[2].partitions[0]], 0) == 5000 assert sum([p.size for p in self.cluster.brokers[2].partitions[1]], 0) == 5000
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 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_skip_topic_different_rf(self): 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.add_replica(b1, 0) partition.add_replica(b2, 1) partition = self.cluster.topics['testTopic3'].partitions[1] partition.add_replica(b1, 0) action = ActionBalanceEven(self.args, self.cluster) assert not action.check_topic_ok(self.cluster.topics['testTopic3'])
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 test_process_cluster_different_rf(self): 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.add_replica(b1, 0) partition.add_replica(b2, 1) partition = self.cluster.topics['testTopic3'].partitions[1] partition.add_replica(b1, 0) self.cluster.topics['testTopic1'].partitions[0].swap_replica_positions(b1, b2) action = ActionBalanceEven(self.args, self.cluster) action.process_cluster() assert self.is_cluster_even(skip_topics=['testTopic3'])
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_odd_count(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] self.cluster.add_topic(Topic("testTopic3", 1)) partition = self.cluster.topics['testTopic3'].partitions[0] partition.add_replica(b1, 0) partition.add_replica(b2, 1) self.cluster.topics['testTopic1'].partitions[0].swap_replica_positions( b1, b2) action = ActionBalanceCount(self.args, self.cluster) action.process_cluster() b1_0 = len(b1.partitions[0]) b1_1 = len(b1.partitions[1]) assert b1_0 == len(b2.partitions[1]) assert b1_1 == len(b2.partitions[0]) assert b1_0 >= 2 and b1_0 <= 3 assert b1_1 >= 2 and b1_1 <= 3
def test_process_cluster_one_move(self): 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.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) self.cluster.topics['testTopic1'].partitions[0].swap_replica_positions( b1, b2) action = ActionBalanceCount(self.args, self.cluster) action.process_cluster() assert len(b1.partitions[0]) == 3 assert len(b1.partitions[1]) == 3 assert len(b2.partitions[0]) == 3 assert len(b2.partitions[1]) == 3
def add_partitions(self, pos, num): topic = Topic('testTopic', num) self.broker.partitions[pos] = [] for i in range(num): self.broker.partitions[pos].append(topic.partitions[i])
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.cluster = Cluster(retention=500000) self.topic = Topic('testTopic', 1) self.topic.cluster = self.cluster self.topic.cluster.topics['testTopic'] = self.topic
class TopicAndPartitionTests(unittest.TestCase): def setUp(self): self.cluster = Cluster(retention=500000) self.topic = Topic('testTopic', 1) self.topic.cluster = self.cluster self.topic.cluster.topics['testTopic'] = self.topic def test_topic_create(self): assert self.topic.name == 'testTopic' assert len(self.topic.partitions) == 1 assert self.topic.cluster == self.cluster def test_partition_create(self): assert len(self.topic.partitions) == 1 assert self.topic.partitions[0].topic == self.topic assert self.topic.partitions[0].num == 0 assert self.topic.partitions[0].replicas == [] assert self.topic.partitions[0].size == 0 def test_topic_equality(self): topic2 = Topic('testTopic', 1) assert self.topic == topic2 def test_topic_equality_with_different_partitions(self): topic2 = Topic('testTopic', 2) assert self.topic == topic2 def test_topic_inequality(self): topic2 = Topic('anotherTestTopic', 1) assert self.topic != topic2 def test_topic_equality_typeerror(self): self.assertRaises(TypeError, self.topic.__eq__, None) def test_topic_copy(self): topic2 = self.topic.copy() assert topic2.name == 'testTopic' assert len(topic2.partitions) == 0 assert self.topic.cluster == topic2.cluster assert self.topic is not topic2 def test_partition_equality(self): partition2 = Partition(self.topic, 0) assert self.topic.partitions[0] == partition2 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_inequality_on_topic_name(self): topic2 = Topic('anotherTestTopic', 1) assert self.topic.partitions[0] != topic2.partitions[0] def test_partition_inequality_on_partition_num(self): partition2 = Partition(self.topic, 1) assert self.topic.partitions[0] != partition2 def test_partition_equality_typeerror(self): self.assertRaises(TypeError, self.topic.partitions[0].__eq__, None) def test_partition_copy_without_replicas(self): partition2 = self.topic.partitions[0].copy() assert self.topic.partitions[0] == partition2 assert self.topic.partitions[0] is not partition2 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 test_add_partition_to_topic(self): partition2 = Partition(self.topic, 1) self.topic.add_partition(partition2) assert len(self.topic.partitions) == 2 assert self.topic.partitions[1].topic == self.topic assert self.topic.partitions[1].num == 1 assert self.topic.partitions[1].replicas == [] assert self.topic.partitions[1].size == 0 def test_set_partition_size(self): self.topic.partitions[0].set_size(100) assert self.topic.partitions[0].size == 100 def test_set_partition_size_skipped(self): self.topic.partitions[0].set_size(100) self.topic.partitions[0].set_size(99) assert self.topic.partitions[0].size == 100 def test_set_partition_size_larger(self): self.topic.partitions[0].set_size(100) self.topic.partitions[0].set_size(200) assert self.topic.partitions[0].size == 200 def test_set_partition_scaled_size_smaller(self): self.topic.retention = 1000000 self.topic.partitions[0].set_size(100) assert self.topic.partitions[0].scaled_size == 50 def test_set_partition_scaled_size_larger(self): self.topic.retention = 100000 self.topic.partitions[0].set_size(100) assert self.topic.partitions[0].scaled_size == 500 def test_partition_dict_for_reassignment_without_replicas(self): expected = {"topic": 'testTopic', "partition": 0, "replicas": []} assert self.topic.partitions[0].dict_for_reassignment() == expected 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 test_partition_dict_for_replica_election(self): expected = {"topic": 'testTopic', "partition": 0} assert self.topic.partitions[0].dict_for_replica_election() == expected
def test_partition_inequality_on_topic_name(self): topic2 = Topic('anotherTestTopic', 1) assert self.topic.partitions[0] != topic2.partitions[0]
def test_topic_inequality(self): topic2 = Topic('anotherTestTopic', 1) assert self.topic != topic2
def add_partitions_to_broker(self, broker_id, pos, num): topic = Topic('testTopic', num) self.cluster.brokers[broker_id].partitions[pos] = [] for i in range(num): self.cluster.brokers[broker_id].partitions[pos].append( topic.partitions[i])
def test_topic_equality(self): topic2 = Topic('testTopic', 1) assert self.topic == topic2
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 setUp(self): self.topic = Topic('testTopic', 1)
class TopicAndPartitionTests(unittest.TestCase): def setUp(self): self.topic = Topic('testTopic', 1) def test_topic_create(self): assert self.topic.name == 'testTopic' assert len(self.topic.partitions) == 1 assert self.topic.cluster is None def test_partition_create(self): assert len(self.topic.partitions) == 1 assert self.topic.partitions[0].topic == self.topic assert self.topic.partitions[0].num == 0 assert self.topic.partitions[0].replicas == [] assert self.topic.partitions[0].size == 0 def test_topic_equality(self): topic2 = Topic('testTopic', 1) assert self.topic == topic2 def test_topic_equality_with_different_partitions(self): topic2 = Topic('testTopic', 2) assert self.topic == topic2 def test_topic_inequality(self): topic2 = Topic('anotherTestTopic', 1) assert self.topic != topic2 def test_topic_equality_typeerror(self): self.assertRaises(TypeError, self.topic.__eq__, None) def test_topic_copy(self): topic2 = self.topic.copy() assert topic2.name == 'testTopic' assert len(topic2.partitions) == 0 assert self.topic.cluster == topic2.cluster assert self.topic is not topic2 def test_partition_equality(self): partition2 = Partition(self.topic, 0) assert self.topic.partitions[0] == partition2 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_inequality_on_topic_name(self): topic2 = Topic('anotherTestTopic', 1) assert self.topic.partitions[0] != topic2.partitions[0] def test_partition_inequality_on_partition_num(self): partition2 = Partition(self.topic, 1) assert self.topic.partitions[0] != partition2 def test_partition_equality_typeerror(self): self.assertRaises(TypeError, self.topic.partitions[0].__eq__, None) def test_partition_copy_without_replicas(self): partition2 = self.topic.partitions[0].copy() assert self.topic.partitions[0] == partition2 assert self.topic.partitions[0] is not partition2 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 test_add_partition_to_topic(self): partition2 = Partition(self.topic, 1) self.topic.add_partition(partition2) assert len(self.topic.partitions) == 2 assert self.topic.partitions[1].topic == self.topic assert self.topic.partitions[1].num == 1 assert self.topic.partitions[1].replicas == [] assert self.topic.partitions[1].size == 0 def test_set_partition_size(self): self.topic.partitions[0].set_size(100) assert self.topic.partitions[0].size == 100 def test_set_partition_size_skipped(self): self.topic.partitions[0].set_size(100) self.topic.partitions[0].set_size(99) assert self.topic.partitions[0].size == 100 def test_set_partition_size_larger(self): self.topic.partitions[0].set_size(100) self.topic.partitions[0].set_size(200) assert self.topic.partitions[0].size == 200 def test_partition_dict_for_reassignment_without_replicas(self): expected = {"topic": 'testTopic', "partition": 0, "replicas": []} assert self.topic.partitions[0].dict_for_reassignment() == expected 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 test_partition_dict_for_replica_election(self): expected = {"topic": 'testTopic', "partition": 0} assert self.topic.partitions[0].dict_for_replica_election() == expected
def test_topic_equality_with_different_partitions(self): topic2 = Topic('testTopic', 2) assert self.topic == topic2
def add_topics(self, num, partition_count=2): for i in range(1, num + 1): topic = Topic("testTopic{0}".format(i), partition_count) self.cluster.add_topic(topic)