def test_update_brokers_from_metadata_update_rack(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'wrongrack'
self.client.cluster.add_broker(broker1)
self.client._update_brokers_from_metadata(self.metadata_response)
assert_cluster_has_brokers(self.client.cluster, self.metadata_response)
def test_update_topics_from_metadata_create(self):
# Don't want to test the broker update code here
self.client.cluster.add_broker(Broker('host1.example.com', id=1, port=8031))
self.client.cluster.add_broker(Broker('host2.example.com', id=101, port=8032))
self.client._update_topics_from_metadata(self.metadata_response)
assert_cluster_has_topics(self.client.cluster, self.metadata_response)
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_update_brokers_from_metadata_update_host(self):
broker1 = Broker('wronghost.example.com', id=1, port=8031)
self.client.cluster.add_broker(broker1)
broker1.close = MagicMock()
self.client._update_brokers_from_metadata(self.metadata_response)
assert_cluster_has_brokers(self.client.cluster, self.metadata_response)
broker1.close.assert_called_once()
def setUp(self):
self.configuration = ClientConfiguration(num_retries=3,
retry_backoff=0.5)
self.mock_sock = MagicMock()
self.broker = Broker('brokerhost1.example.com',
id=1,
sock=self.mock_sock,
configuration=self.configuration)
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_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_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_check_partition_swappable_racks_ok(self):
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = Broker('brokerhost3.example.com', id=3)
b4 = Broker('brokerhost4.example.com', id=4)
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 test_try_pick_new_broker(self):
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = Broker('brokerhost3.example.com', id=3)
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 test_send_group_aware_request(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
self.client.cluster.add_broker(broker1)
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.group_coordinator
self.client._send_group_aware_request('testgroup', 'fakerequest')
assert 'testgroup' in self.client.cluster.groups
assert self.client.cluster.groups['testgroup'].coordinator == broker1
broker1.send.assert_called_once_with('fakerequest')
def test_send_group_aware_request(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
self.client.cluster.add_broker(broker1)
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.group_coordinator
self.client._send_group_aware_request('testgroup', 'fakerequest')
assert 'testgroup' in self.client.cluster.groups
assert self.client.cluster.groups['testgroup'].coordinator == broker1
broker1.send.assert_called_once_with('fakerequest')
def test_try_pick_new_broker(self):
b1 = self.cluster.brokers[1]
b2 = self.cluster.brokers[2]
b3 = Broker('brokerhost3.example.com', id=3)
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_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):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.client._connected = True
# Two brokers for the client
broker = Broker('host1.example.com', id=1, port=8031)
broker.rack = 'rack1'
self.client.cluster.add_broker(broker)
broker = Broker('host2.example.com', id=101, port=8032)
broker.rack = 'rack1'
self.client.cluster.add_broker(broker)
self.metadata = topic_metadata()
self.metadata_error = topic_metadata_error()
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 add_brokers_from_zk(cluster, zk):
for b in zk.get_children("/brokers/ids"):
broker_data, bstat = zk.get("/brokers/ids/{0}".format(b))
cluster.add_broker(Broker.create_from_json(int(b), broker_data))
if cluster.num_brokers() == 0:
raise ZookeeperException(
"The cluster specified does not have any brokers")
def setUp(self):
self.broker = Broker('testhost', port=3945)
self.request_classes = {'testklass': {1: TestKlass1}, 'testklassy': {1: TestKlassy1}}
self.request_cmds = {'testklass': TestKlass1,
'testklassv1': TestKlass1,
'testklassy': TestKlassy1,
'testklassyv1': TestKlassy1}
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 main():
# Set up and parse all CLI arguments
args = set_up_arguments()
# Connect to the specified Kafka broker
broker = Broker(args.broker, port=args.port)
broker.connect()
# Loop on reading a command
_cli_loop(broker)
# Disconnect
broker.close()
# Nothing to see here, move along
return 0
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_send_group_aware_request_new_broker(self, mock_broker_class):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
mock_broker_class.return_value = broker1
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.group_coordinator
self.client._send_group_aware_request('testgroup', 'fakerequest')
mock_broker_class.assert_called_once_with('host1.example.com', id=1, port=8031, configuration=self.client.configuration)
assert 1 in self.client.cluster.brokers
assert self.client.cluster.brokers[1] == broker1
assert self.client.cluster.groups['testgroup'].coordinator == broker1
broker1.send.assert_called_once_with('fakerequest')
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_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_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('brokerhost3.example.com', id=3)
b4 = Broker('brokerhost4.example.com', id=4)
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_broker_create_from_json_extended(self):
jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}'
broker2 = Broker.create_from_json(1, jsonstr)
assert broker2.jmx_port == -1
assert broker2.timestamp == "1466985807242"
assert broker2.endpoints == ["PLAINTEXT://10.0.0.10:9092"]
assert broker2.version == 3
assert broker2.port == 9092
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_update_topics_from_metadata_missing_broker(self):
# Don't want to test the broker update code here
self.client.cluster.add_broker(Broker('host1.example.com', id=1, port=8031))
self.client._update_topics_from_metadata(self.metadata_response)
assert_cluster_has_topics(self.client.cluster, self.metadata_response)
assert 101 in self.client.cluster.brokers
assert self.client.cluster.brokers[101].endpoint.hostname is None
def test_broker_get_endpoint(self):
jsonstr = ('{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092", "SSL://10.0.0.10:2834"],'
'"host":"10.0.0.10","version":3,"port":9092}')
broker = Broker.create_from_json(1, jsonstr)
endpoint = broker.get_endpoint("SSL")
assert endpoint.protocol == "SSL"
assert endpoint.hostname == "10.0.0.10"
assert endpoint.port == 2834
def test_add_or_update_group(self):
broker = Broker('host1.example.com', id=1, port=8031)
self.client.cluster.add_broker(broker)
list_group = list_groups()
self.client._add_or_update_group(list_group['groups'][0], 1)
assert 'group1' in self.client.cluster.groups
assert self.client.cluster.groups['group1'].coordinator == broker
assert self.client.cluster.groups['group1'].protocol_type == 'protocol1'
def test_broker_get_endpoint(self):
jsonstr = (
'{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092", "SSL://10.0.0.10:2834"],'
'"host":"10.0.0.10","version":3,"port":9092}')
broker = Broker.create_from_json(1, jsonstr)
endpoint = broker.get_endpoint("SSL")
assert endpoint.protocol == "SSL"
assert endpoint.hostname == "10.0.0.10"
assert endpoint.port == 2834
def _update_or_add_partition(self, partition_metadata, partition):
partition.leader = self.cluster.brokers[partition_metadata['leader']]
for i, replica in enumerate(partition_metadata['replicas']):
if replica not in self.cluster.brokers:
# We have a replica ID that is not a known broker. This can happen if a broker is offline, or
# if the partition is otherwise assigned to a non-existent broker ID. In this case, we need to
# create a broker object for this with no endpoint information as a placeholder.
self.cluster.add_broker(Broker(hostname=None, id=replica))
partition.add_or_update_replica(i, self.cluster.brokers[replica])
partition.delete_replicas(len(partition_metadata['replicas']))
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_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_broker_create_from_json_extended(self):
jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}'
broker2 = Broker.create_from_json(1, jsonstr)
assert broker2.jmx_port == -1
assert broker2.timestamp == "1466985807242"
assert len(broker2.endpoints) == 1
assert broker2.endpoints[0].protocol == 'PLAINTEXT'
assert broker2.endpoints[0].hostname == '10.0.0.10'
assert broker2.endpoints[0].port == 9092
assert broker2.version == 3
assert broker2.port == 9092
def test_send_group_aware_request_new_broker(self, mock_broker_class):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
mock_broker_class.return_value = broker1
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.group_coordinator
self.client._send_group_aware_request('testgroup', 'fakerequest')
mock_broker_class.assert_called_once_with(
'host1.example.com',
id=1,
port=8031,
configuration=self.client.configuration)
assert 1 in self.client.cluster.brokers
assert self.client.cluster.brokers[1] == broker1
assert self.client.cluster.groups['testgroup'].coordinator == broker1
broker1.send.assert_called_once_with('fakerequest')
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_send_any_broker(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
broker2 = Broker('host2.example.com', id=101, port=8032)
broker2.rack = 'rack1'
broker2.send = MagicMock()
broker2.send.return_value = (2, 'otherresponse')
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
val = self.client._send_any_broker('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
assert val == 'otherresponse'
def test_close(self):
# Two brokers for the client
broker = Broker('host1.example.com', id=1, port=8031)
broker.rack = 'rack1'
broker.close = MagicMock()
self.client.cluster.add_broker(broker)
broker = Broker('host2.example.com', id=101, port=8032)
broker.rack = 'rack1'
broker.close = MagicMock()
self.client.cluster.add_broker(broker)
self.client.close()
for broker_id in self.client.cluster.brokers:
self.client.cluster.brokers[broker_id].close.assert_called_once()
def test_send_any_broker(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
broker2 = Broker('host2.example.com', id=101, port=8032)
broker2.rack = 'rack1'
broker2.send = MagicMock()
broker2.send.return_value = (2, 'otherresponse')
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
val = self.client._send_any_broker('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
assert val == 'otherresponse'
def test_send_all_error(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.side_effect = ConnectionError
broker2 = Broker('host2.example.com', id=101, port=8032)
broker2.rack = 'rack1'
broker2.send = MagicMock()
broker2.send.side_effect = ConnectionError
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
self.assertRaises(ConnectionError, self.client._send_any_broker, 'fakerequest')
broker1.send.assert_called_once_with('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
def test_close(self):
# Two brokers for the client
broker = Broker('host1.example.com', id=1, port=8031)
broker.rack = 'rack1'
broker.close = MagicMock()
self.client.cluster.add_broker(broker)
broker = Broker('host2.example.com', id=101, port=8032)
broker.rack = 'rack1'
broker.close = MagicMock()
self.client.cluster.add_broker(broker)
self.client.close()
for broker_id in self.client.cluster.brokers:
self.client.cluster.brokers[broker_id].close.assert_called_once()
def test_send_all_brokers_error(self):
broker1 = Broker('host1.example.com', id=1, port=8031)
broker1.rack = 'rack1'
broker1.send = MagicMock()
broker1.send.return_value = (1, 'fakeresponse')
broker2 = Broker('host2.example.com', id=101, port=8032)
broker2.rack = 'rack1'
broker2.send = MagicMock()
broker2.send.side_effect = ConnectionError
self.client.cluster.add_broker(broker1)
self.client.cluster.add_broker(broker2)
val = self.client._send_all_brokers('fakerequest')
broker1.send.assert_called_once_with('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
assert val[1] == 'fakeresponse'
assert val[101] is None
def _maybe_bootstrap_cluster(self, broker_port):
"""Attempt to bootstrap the cluster information using the given broker"""
broker = Broker(broker_port[0], port=broker_port[1], configuration=self.configuration)
try:
broker.connect()
except ConnectionError:
# Just skip to the next bootstrap broker
return False
# Fetch topic metadata for all topics/brokers
req = TopicMetadataV1Request({'topics': None})
correlation_id, metadata = broker.send(req)
broker.close()
# Add brokers and topics to cluster
self._controller_id = metadata['controller_id']
self._update_from_metadata(metadata)
self._last_full_metadata = time.time()
return True
class BrokerTests(unittest.TestCase):
def setUp(self):
self.configuration = ClientConfiguration(num_retries=3, retry_backoff=0.5)
self.mock_sock = MagicMock()
self.broker = Broker('brokerhost1.example.com', id=1, sock=self.mock_sock, configuration=self.configuration)
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 test_broker_create(self):
assert self.broker.id == 1
assert self.broker.hostname == 'brokerhost1.example.com'
assert self.broker.partitions == {}
assert isinstance(self.broker._sock, MagicMock)
def test_broker_create_from_json_basic(self):
jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}'
broker2 = Broker.create_from_json(1, jsonstr)
assert broker2.id == 1
assert broker2.hostname == '10.0.0.10'
assert broker2.partitions == {}
def test_broker_get_endpoint(self):
jsonstr = ('{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092", "SSL://10.0.0.10:2834"],'
'"host":"10.0.0.10","version":3,"port":9092}')
broker = Broker.create_from_json(1, jsonstr)
endpoint = broker.get_endpoint("SSL")
assert endpoint.protocol == "SSL"
assert endpoint.hostname == "10.0.0.10"
assert endpoint.port == 2834
def test_broker_create_from_json_extended(self):
jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}'
broker2 = Broker.create_from_json(1, jsonstr)
assert broker2.jmx_port == -1
assert broker2.timestamp == "1466985807242"
assert len(broker2.endpoints) == 1
assert broker2.endpoints[0].protocol == 'PLAINTEXT'
assert broker2.endpoints[0].hostname == '10.0.0.10'
assert broker2.endpoints[0].port == 9092
assert broker2.version == 3
assert broker2.port == 9092
def test_broker_create_from_json_bad_host(self):
jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"hostname":"10.0.0.10","version":3,"port":9092}'
self.assertRaises(ConfigurationException, Broker.create_from_json, 1, jsonstr)
def test_broker_create_from_json_bad_jmx_port(self):
jsonstr = '{"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}'
broker2 = Broker.create_from_json(1, jsonstr)
assert broker2.hostname == '10.0.0.10'
def test_broker_copy_without_partitions(self):
broker2 = self.broker.copy()
assert broker2.id == 1
assert broker2.hostname == 'brokerhost1.example.com'
assert self.broker is not broker2
def test_broker_copy_with_partitions(self):
self.add_partitions(0, 1)
broker2 = self.broker.copy()
assert broker2.id == 1
assert broker2.hostname == 'brokerhost1.example.com'
assert broker2.partitions == {}
assert self.broker is not broker2
def test_broker_num_leaders(self):
self.add_partitions(0, 2)
assert self.broker.num_leaders() == 2
def test_broker_num_partitions_at_position_zero(self):
self.add_partitions(1, 2)
assert self.broker.num_partitions_at_position(0) == 0
def test_broker_num_partitions_single_position(self):
self.add_partitions(0, 2)
assert self.broker.num_partitions() == 2
def test_broker_num_partitions_two_position(self):
self.add_partitions(0, 2)
self.add_partitions(1, 1)
assert self.broker.num_partitions() == 3
def test_broker_percent_leaders(self):
self.add_partitions(0, 2)
self.add_partitions(1, 2)
assert self.broker.percent_leaders() == 50.0
def test_broker_percent_leaders_zero(self):
self.add_partitions(1, 1)
assert self.broker.percent_leaders() == 0.0
def test_broker_percent_leaders_no_partitions(self):
assert self.broker.percent_leaders() == 0.0
def test_broker_equality(self):
broker2 = Broker('brokerhost1.example.com', id=1)
assert self.broker == broker2
def test_broker_inequality_hostname(self):
broker2 = Broker('brokerhost2.example.com', id=1)
assert self.broker != broker2
def test_broker_inequality_id(self):
broker2 = Broker('brokerhost1.example.com', id=2)
assert self.broker != broker2
def test_broker_equality_typeerror(self):
self.assertRaises(TypeError, self.broker.__eq__, None)
def test_broker_get_socket(self):
sock = self.broker._get_socket(None)
assert isinstance(sock, socket.socket)
@patch('kafka.tools.models.broker.socket.socket')
def test_broker_get_socket_wrapped(self, mock_socket):
mock_context = MagicMock()
mock_context.wrap_socket.return_value = 'fakewrappedsocket'
mock_socket.return_value = 'fakesocket'
sock = self.broker._get_socket(mock_context)
mock_context.wrap_socket.assert_called_once_with('fakesocket', server_hostname='brokerhost1.example.com')
print(sock)
assert sock == 'fakewrappedsocket'
def test_broker_connect(self):
self.broker.connect()
self.mock_sock.connect.assert_called_once_with(('brokerhost1.example.com', 9092))
def test_broker_connect_error(self):
self.mock_sock.connect.side_effect = socket.error
self.assertRaises(ConnectionError, self.broker.connect)
def test_broker_close(self):
self.broker.close()
self.mock_sock.shutdown.assert_called_once()
self.mock_sock.close.assert_called_once()
def test_broker_close_not_open(self):
self.mock_sock.shutdown.side_effect = OSError()
self.broker.close()
self.mock_sock.close.assert_called_once()
@patch.object(Broker, '_read_bytes')
def test_broker_single_send(self, mock_read_bytes):
# recv is called to get the response size (first 4 bytes)
self.mock_sock.recv.return_value = b'\x00\x00\x00\x10'
# read_bytes returns the response payload, minus the response size
# correlation_id (4), error (2), array of (api_key (2), min_version (2), max_version (2))
mock_read_bytes.return_value = b'\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x01\x01\x01\x02\x02'
request = ApiVersionsV0Request({})
(correlation_id, response) = self.broker._single_send(request)
# Check that the request was encoded properly
self.mock_sock.sendall.assert_called_once_with(bytearray(b'\x00\x00\x00\x15\x00\x12\x00\x00\x00\x00\x00\x01\x00\x0bkafka-tools'))
self.mock_sock.recv.assert_called_once_with(4)
mock_read_bytes.assert_called_once_with(16)
# Check that the response is what was expected to be decoded
assert isinstance(response, ApiVersionsV0Response)
assert response.correlation_id == 1
assert response['error'] == 0
assert isinstance(response['api_versions'], collections.Sequence)
assert len(response['api_versions']) == 1
assert len(response['api_versions'][0]) == 3
assert response['api_versions'][0]['api_key'] == 1
assert response['api_versions'][0]['min_version'] == 257
assert response['api_versions'][0]['max_version'] == 514
# Correlation ID must be incremented after each request
assert self.broker._correlation_id == 2
@patch.object(Broker, '_read_bytes')
def test_broker_single_send_error(self, mock_read_bytes):
self.mock_sock.recv.side_effect = socket.error
mock_read_bytes.return_value = b'\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x01\x01\x01\x02\x02'
request = ApiVersionsV0Request({})
self.assertRaises(ConnectionError, self.broker._single_send, request)
@patch.object(ByteBuffer, 'getInt32')
def test_broker_single_send_short(self, mock_getint):
self.mock_sock.recv.return_value = b'\x00\x00\x00\x10'
mock_getint.side_effect = EOFError
request = ApiVersionsV0Request({})
self.assertRaises(ConnectionError, self.broker._single_send, request)
@patch.object(Broker, '_single_send')
@patch.object(Broker, 'connect')
@patch.object(Broker, 'close')
def test_broker_send(self, mock_close, mock_connect, mock_send):
mock_send.return_value = 'fakeresponse'
self.broker.send('fakerequest')
mock_connect.assert_not_called()
mock_close.assert_not_called()
mock_send.assert_called_once_with('fakerequest')
@patch.object(Broker, '_single_send')
@patch.object(Broker, 'connect')
@patch.object(Broker, 'close')
def test_broker_send_retry(self, mock_close, mock_connect, mock_send):
def close_broker():
self.broker._sock = None
mock_send.side_effect = [ConnectionError, 'fakeresponse']
mock_close.side_effect = close_broker
self.broker.send('fakerequest')
mock_close.assert_called_once_with()
mock_connect.assert_called_once_with()
mock_send.assert_has_calls([call('fakerequest'), call('fakerequest')])
@patch.object(Broker, '_single_send')
@patch.object(Broker, 'connect')
@patch.object(Broker, 'close')
def test_broker_send_exhaust_retry(self, mock_close, mock_connect, mock_send):
def close_broker():
self.broker._sock = None
mock_send.side_effect = ConnectionError
mock_close.side_effect = close_broker
self.assertRaises(ConnectionError, self.broker.send, 'fakerequest')
mock_close.assert_has_calls([call(), call()])
mock_connect.assert_has_calls([call(), call()])
mock_send.assert_has_calls([call('fakerequest'), call('fakerequest'), call('fakerequest')])
def test_broker_read_bytes(self):
self.mock_sock.recv.return_value = b'\x01\x02\x03\x04'
response = self.broker._read_bytes(4)
self.mock_sock.recv.assert_called_once_with(4)
assert response == b'\x01\x02\x03\x04'
def test_broker_read_bytes_concat(self):
self.mock_sock.recv.side_effect = [b'\x01\x02', b'\x03\x04']
response = self.broker._read_bytes(4)
self.mock_sock.recv.assert_has_calls([call(4), call(2)])
assert response == b'\x01\x02\x03\x04'
def test_broker_read_bytes_nodata(self):
self.mock_sock.recv.return_value = b''
self.assertRaises(socket.error, self.broker._read_bytes, 4)
def test_broker_read_bytes_error(self):
self.mock_sock.recv.side_effect = socket.error
self.assertRaises(socket.error, self.broker._read_bytes, 4)
def setUp(self):
self.configuration = ClientConfiguration(num_retries=3, retry_backoff=0.5)
self.mock_sock = MagicMock()
self.broker = Broker('brokerhost1.example.com', id=1, sock=self.mock_sock, configuration=self.configuration)
def test_broker_create_from_json_bad_jmx_port(self):
jsonstr = '{"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}'
broker2 = Broker.create_from_json(1, jsonstr)
assert broker2.hostname == '10.0.0.10'
