def setUp(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker("brokerhost1.example.com", id=1))
self.cluster.add_broker(Broker("brokerhost2.example.com", id=2))
self.cluster.add_broker(Broker("brokerhost3.example.com", id=3))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
def test_assure_topic_has_partitions_add_one(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
assert len(topic.partitions) == 1
topic.assure_has_partitions(2)
assert len(topic.partitions) == 2
assert topic.partitions[1].num == 1
def test_assure_topic_has_partitions_add_one(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
assert len(topic.partitions) == 1
topic.assure_has_partitions(2)
assert len(topic.partitions) == 2
assert topic.partitions[1].num == 1
def create_cluster_onehost(self):
cluster = Cluster()
cluster.add_broker(Broker("brokerhost1.example.com", id=1))
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 set_up_cluster_4broker():
cluster = Cluster()
cluster.add_broker(Broker("brokerhost1.example.com", id=1))
cluster.add_broker(Broker("brokerhost2.example.com", id=2))
cluster.add_broker(Broker("brokerhost3.example.com", id=3))
cluster.add_broker(Broker("brokerhost4.example.com", id=4))
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 setUp(self):
self.topic = Topic('testTopic', 10)
self.broker = Broker('brokerhost1.example.com', id=1)
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_set_offsets_from_fetch_error(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = offset_fetch_error()
self.assertRaises(OffsetError, offsets.set_offsets_from_fetch,
response['responses'][0]['partition_responses'])
def test_topic_offsets_create(self):
topic = Topic('topic1', 3)
offsets = TopicOffsets(topic)
assert offsets.topic == topic
assert len(offsets.partitions) == 3
for partition in offsets.partitions:
assert partition == -1
def test_process_cluster_empty_one(self):
self.cluster.add_broker(Broker('brokerhost3.example.com', id=3))
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('brokerhost3.example.com', id=3))
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 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_get_topic_expired(self):
topic = Topic('topic1', 1)
topic._last_updated = 100
self.client.cluster.add_topic(topic)
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.metadata
val = self.client.get_topic('topic1')
self.client._send_any_broker.assert_called_once()
assert isinstance(self.client._send_any_broker.call_args[0][0], TopicMetadataV1Request)
assert isinstance(val, Topic)
assert val.name == 'topic1'
assert len(val.partitions) == 2
def test_process_cluster_zero_partitions(self):
self.cluster.add_broker(Broker("brokerhost3.example.com", id=3))
self.cluster.brokers[1].rack = "c"
self.cluster.add_topic(Topic("testTopic3", 1))
partition = self.cluster.topics['testTopic3'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition.add_replica(self.cluster.brokers[3], 2)
self.cluster.topics['testTopic3'].partitions[0].set_size(2000)
self.args.exclude_topics = ['testTopic3']
action = ActionBalanceRate(self.args, self.cluster)
action.process_cluster()
assert sum(
[p.scaled_size for p in self.cluster.brokers[1].partitions[0]],
0) == 3000
assert sum(
[p.scaled_size for p in self.cluster.brokers[1].partitions[1]],
0) == 2000
assert sum(
[p.scaled_size for p in self.cluster.brokers[2].partitions[0]],
0) == 3000
assert sum(
[p.scaled_size for p in self.cluster.brokers[2].partitions[1]],
0) == 4000
assert sum(
[p.scaled_size for p in self.cluster.brokers[3].partitions[0]],
0) == 2000
assert sum(
[p.scaled_size for p in self.cluster.brokers[3].partitions[1]],
0) == 2000
assert sum(
[p.scaled_size for p in self.cluster.brokers[3].partitions[2]],
0) == 2000
def test_list_topics(self):
self.client.cluster.add_topic(Topic('topic1', 1))
self.client._maybe_update_full_metadata = MagicMock()
topics = self.client.list_topics()
self.client._maybe_update_full_metadata.assert_called_once_with(True)
assert topics == ['topic1']
def test_set_offsets_from_fetch(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = offset_fetch()
offsets.set_offsets_from_fetch(
response['responses'][0]['partition_responses'])
assert offsets.partitions[0] == 4829
assert offsets.partitions[1] == 8904
def test_set_offsets_empty_from_list(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = list_offset_none()
offsets.set_offsets_from_list(
response['responses'][0]['partition_responses'])
assert offsets.partitions[0] == 4829
assert offsets.partitions[1] == -1
def test_get_topic_error(self):
topic = Topic('topic1', 1)
self.client.cluster.add_topic(topic)
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.metadata_error
self.assertRaises(TopicError, self.client.get_topic, 'topic1', cache=False)
self.client._send_any_broker.assert_called_once()
assert isinstance(self.client._send_any_broker.call_args[0][0], TopicMetadataV1Request)
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(None, id=replica))
newtopic.partitions[int(partition)].add_replica(cluster.brokers[replica], i)
cluster.add_topic(newtopic)
def test_delete_replicas_from_partition_all(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
partition = topic.partitions[0]
assert len(partition.replicas) == 2
partition.delete_replicas(0)
assert len(partition.replicas) == 0
def test_maybe_delete_topics_not_in_metadata(self):
# Don't want to test the broker update code here
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
self.client.cluster.add_topic(topic)
topic.partitions[0].add_replica(broker2)
topic.partitions[0].add_replica(broker1)
topic = Topic('topic2', 1)
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
self.client.cluster.add_topic(topic)
topic.partitions[0].add_replica(broker2)
topic.partitions[0].add_replica(broker1)
self.client._maybe_delete_topics_not_in_metadata(self.metadata_response, delete=True)
assert 'topic2' not in self.client.cluster.topics
assert 'topic1' in self.client.cluster.topics
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("brokerhost1.example.com", id=1))
cluster.add_broker(Broker("brokerhost2.example.com", id=2))
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_get_topic_existing_cached(self):
topic = Topic('topic1', 2)
self.client.cluster.add_topic(topic)
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.metadata
self.client._update_from_metadata = MagicMock()
val = self.client.get_topic('topic1')
self.client._send_any_broker.assert_not_called()
self.client._update_from_metadata.assert_not_called()
assert isinstance(val, Topic)
assert val.name == 'topic1'
def set_up_cluster():
cluster = Cluster()
cluster.add_broker(Broker("brokerhost1.example.com", id=1))
cluster.add_broker(Broker("brokerhost2.example.com", id=2))
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_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 test_update_topics_from_metadata_update_replicas(self):
# Don't want to test the broker update code here
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
broker3 = Broker('host3.example.com', id=304, port=8033)
topic = Topic('topic1', 2)
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
self.client.cluster.add_broker(broker3)
self.client.cluster.add_topic(topic)
topic.partitions[0].add_replica(broker3)
topic.partitions[0].add_replica(broker1)
topic.partitions[1].add_replica(broker2)
topic.partitions[1].add_replica(broker1)
self.client._update_topics_from_metadata(self.metadata_response)
assert_cluster_has_topics(self.client.cluster, self.metadata_response)
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 _update_topics_from_metadata(self, metadata, delete=False):
"""
Given a Metadata response (either V0 or V1 will work), update the topic information
for this cluster.
Args:
metadata (MetadataV1Response): A metadata response to create or update topics for
delete (boolean): If True, delete topics from the cluster that are not present in the
metadata response
Raises:
IndexError: If the brokers in the metadata object are not defined in the cluster
"""
for t in metadata['topics']:
if t['name'] not in self.cluster.topics:
self.cluster.add_topic(Topic(t['name'], len(t['partitions'])))
topic = self.cluster.topics[t['name']]
topic._last_updated = time.time()
topic.assure_has_partitions(len(t['partitions']))
for p in t['partitions']:
self._update_or_add_partition(p, topic.partitions[p['id']])
self._maybe_delete_topics_not_in_metadata(metadata, delete)
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_updated_since(self):
self.topic._last_updated = 100
assert self.topic.updated_since(99)
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('testhost1', id=1)
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('testhost1', id=1)
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('testhost1', id=1)
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_add_or_update_replica_nochange(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
partition = Partition('topic1', 0)
partition.add_replica(broker1)
partition.add_replica(broker2)
partition.add_or_update_replica(0, broker1)
assert partition.replicas[0] == broker1
def test_add_or_update_replica_new(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
partition = Partition('topic1', 0)
partition.add_replica(broker1)
partition.add_replica(broker2)
partition.add_or_update_replica(2, broker1)
assert partition.replicas[2] == broker1
def test_add_or_update_replica_swap(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
partition = Partition('topic1', 0)
partition.add_replica(broker1)
partition.add_replica(broker2)
partition.add_or_update_replica(1, broker1)
assert partition.replicas[1] == broker1
def test_assure_topic_has_partitions_nochange(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 2)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
topic.partitions[1].add_replica(broker2)
topic.partitions[1].add_replica(broker1)
assert len(topic.partitions) == 2
topic.assure_has_partitions(2)
assert len(topic.partitions) == 2
def test_assure_topic_has_partitions_one(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 2)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
topic.partitions[1].add_replica(broker2)
topic.partitions[1].add_replica(broker1)
assert len(topic.partitions) == 2
topic.assure_has_partitions(1)
assert len(topic.partitions) == 1
assert topic.partitions[0].num == 0
def test_assure_topic_has_partitions_add_one(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
assert len(topic.partitions) == 1
topic.assure_has_partitions(2)
assert len(topic.partitions) == 2
assert topic.partitions[1].num == 1
def test_assure_topic_has_partitions_all(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 2)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
topic.partitions[1].add_replica(broker2)
topic.partitions[1].add_replica(broker1)
assert len(topic.partitions) == 2
topic.assure_has_partitions(0)
assert len(topic.partitions) == 0
def test_delete_replicas_from_partition_nochange(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
partition = topic.partitions[0]
assert len(partition.replicas) == 2
partition.delete_replicas(2)
assert len(partition.replicas) == 2
def test_delete_replicas_from_partition_one(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
partition = topic.partitions[0]
assert len(partition.replicas) == 2
partition.delete_replicas(1)
assert len(partition.replicas) == 1
assert partition.replicas[0] == broker1
def test_delete_replicas_from_partition_all(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker2 = Broker('host2.example.com', id=101, port=8032)
topic = Topic('topic1', 1)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker2)
partition = topic.partitions[0]
assert len(partition.replicas) == 2
partition.delete_replicas(0)
assert len(partition.replicas) == 0
def test_topic_offsets_create(self):
topic = Topic('topic1', 3)
offsets = TopicOffsets(topic)
assert offsets.topic == topic
assert len(offsets.partitions) == 3
for partition in offsets.partitions:
assert partition == -1
def test_set_offsets_from_list(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = list_offset()
offsets.set_offsets_from_list(response['responses'][0]['partition_responses'])
assert offsets.partitions[0] == 4829
assert offsets.partitions[1] == 8904
def test_set_offsets_empty_from_list(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = list_offset_none()
offsets.set_offsets_from_list(response['responses'][0]['partition_responses'])
assert offsets.partitions[0] == 4829
assert offsets.partitions[1] == -1
def test_set_offsets_from_list_error(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = list_offset_error()
self.assertRaises(OffsetError, offsets.set_offsets_from_list, response['responses'][0]['partition_responses'])
def test_set_offsets_from_fetch(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = offset_fetch()
offsets.set_offsets_from_fetch(response['responses'][0]['partition_responses'])
assert offsets.partitions[0] == 4829
assert offsets.partitions[1] == 8904
def test_set_offsets_from_fetch_error(self):
topic = Topic('topic1', 2)
offsets = TopicOffsets(topic)
response = offset_fetch_error()
self.assertRaises(OffsetError, offsets.set_offsets_from_fetch, response['responses'][0]['partition_responses'])
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 setUp(self):
self.cluster = Cluster(retention=500000)
self.topic = Topic('testTopic', 1)
self.topic.cluster = self.cluster
self.topic.cluster.topics['testTopic'] = self.topic
