def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
# Get the broker and topic from a metadata update
self.client._update_from_metadata(topic_metadata())
self.list_offset = list_offset()
self.list_offset_error = list_offset_error()
self.list_offset_request = {
1: {
'replica_id':
-1,
'topics': [{
'topic':
'topic1',
'partitions': [{
'partition': 0,
'timestamp': Client.OFFSET_LATEST,
'max_num_offsets': 1
}, {
'partition': 1,
'timestamp': Client.OFFSET_EARLIEST,
'max_num_offsets': 1
}]
}]
}
}
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.client._connected = True
# Get the broker and topic from a metadata update
self.client._update_from_metadata(topic_metadata())
self.group = Group('testgroup')
self.offset_fetch = offset_fetch()
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 setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
# Get the broker and topic from a metadata update
self.client._update_from_metadata(topic_metadata())
self.list_offset = list_offset()
self.list_offset_error = list_offset_error()
self.list_offset_request = {1: {'replica_id': -1,
'topics': [{'topic': 'topic1',
'partitions': [{'partition': 0,
'timestamp': Client.OFFSET_LATEST,
'max_num_offsets': 1},
{'partition': 1,
'timestamp': Client.OFFSET_EARLIEST,
'max_num_offsets': 1}]}]}}
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.group_coordinator = group_coordinator()
self.coordinator_error = group_coordinator_error()
class SendHelperTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.group_coordinator = group_coordinator()
self.coordinator_error = group_coordinator_error()
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')
@patch('kafka.tools.client.Broker')
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_send_group_aware_request_error(self):
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.coordinator_error
self.assertRaises(GroupError, self.client._send_group_aware_request, 'testgroup', 'fakerequest')
def test_send_all_brokers(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_all_brokers('fakerequest')
broker1.send.assert_called_once_with('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
assert val[1] == 'fakeresponse'
assert val[101] == 'otherresponse'
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
@patch('kafka.tools.client.shuffle', lambda x: sorted(x))
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'
@patch('kafka.tools.client.shuffle', lambda x: sorted(x))
def test_send_any_broker_one_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_any_broker('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
broker1.send.assert_called_once_with('fakerequest')
assert val == 'fakeresponse'
@patch('kafka.tools.client.shuffle', lambda x: sorted(x))
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_raise_if_not_connected(self):
self.assertRaises(ConnectionError, self.client._raise_if_not_connected)
def test_create_config_zkconnect(self):
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
assert test_client.configuration.zkconnect == 'zk.example.com:2181/kafka-cluster'
class OffsetTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
# Get the broker and topic from a metadata update
self.client._update_from_metadata(topic_metadata())
self.list_offset = list_offset()
self.list_offset_error = list_offset_error()
self.list_offset_request = {1: {'replica_id': -1,
'topics': [{'topic': 'topic1',
'partitions': [{'partition': 0,
'timestamp': Client.OFFSET_LATEST,
'max_num_offsets': 1},
{'partition': 1,
'timestamp': Client.OFFSET_EARLIEST,
'max_num_offsets': 1}]}]}}
def test_send_list_offsets_to_brokers(self):
self.client.cluster.brokers[1].send = MagicMock()
self.client.cluster.brokers[1].send.return_value = (1, self.list_offset)
val = self.client._send_list_offsets_to_brokers(self.list_offset_request)
self.client.cluster.brokers[1].send.assert_called_once()
req = self.client.cluster.brokers[1].send.call_args[0][0]
assert req['replica_id'] == -1
assert len(req['topics']) == 1
assert req['topics'][0]['topic'] == 'topic1'
assert len(req['topics'][0]['partitions']) == 2
assert req['topics'][0]['partitions'][0]['partition'] == 0
assert req['topics'][0]['partitions'][0]['timestamp'] == Client.OFFSET_LATEST
assert req['topics'][0]['partitions'][0]['max_num_offsets'] == 1
assert req['topics'][0]['partitions'][1]['partition'] == 1
assert req['topics'][0]['partitions'][1]['timestamp'] == Client.OFFSET_EARLIEST
assert req['topics'][0]['partitions'][1]['max_num_offsets'] == 1
assert len(val) == 1
assert 'topic1' in val
assert isinstance(val['topic1'], TopicOffsets)
assert val['topic1'].topic == self.client.cluster.topics['topic1']
assert len(val['topic1'].partitions) == 2
assert val['topic1'].partitions[0] == 4829
assert val['topic1'].partitions[1] == 8904
def test_send_list_offsets_to_brokers_connection_error(self):
self.client.cluster.brokers[1].send = MagicMock()
self.client.cluster.brokers[1].send.side_effect = ConnectionError
self.assertRaises(ConnectionError, self.client._send_list_offsets_to_brokers, self.list_offset_request)
def test_send_list_offsets_to_brokers_offsets_error(self):
self.client.cluster.brokers[1].send = MagicMock()
self.client.cluster.brokers[1].send.return_value = (1, self.list_offset_error)
self.assertRaises(OffsetError, self.client._send_list_offsets_to_brokers, self.list_offset_request)
def test_send_set_offset_request_bad_offsets(self):
self.assertRaises(TypeError, self.client._send_set_offset_request, 'testgroup', ['notatopicoffsets'])
def test_send_set_offset_request_bad_topic(self):
offsets = TopicOffsets(self.client.cluster.topics['topic1'])
offsets.topic = None
self.assertRaises(TypeError, self.client._send_set_offset_request, 'testgroup', [offsets])
def test_send_set_offset_request(self):
offsets = TopicOffsets(self.client.cluster.topics['topic1'])
offsets.partitions[0] = 2342
offsets.partitions[1] = 8793
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.return_value = 'responseobject'
val = self.client._send_set_offset_request('testgroup', [offsets])
assert val == 'responseobject'
self.client._send_group_aware_request.assert_called_once()
assert self.client._send_group_aware_request.call_args[0][0] == 'testgroup'
req = self.client._send_group_aware_request.call_args[0][1]
assert isinstance(req, OffsetCommitV2Request)
assert req['group_id'] == 'testgroup'
assert req['group_generation_id'] == -1
assert req['member_id'] == ''
assert req['retention_time'] == -1
assert len(req['topics']) == 1
assert req['topics'][0]['topic'] == 'topic1'
assert len(req['topics'][0]['partitions']) == 2
assert req['topics'][0]['partitions'][0]['partition'] == 0
assert req['topics'][0]['partitions'][0]['offset'] == 2342
assert req['topics'][0]['partitions'][0]['metadata'] is None
assert req['topics'][0]['partitions'][1]['partition'] == 1
assert req['topics'][0]['partitions'][1]['offset'] == 8793
assert req['topics'][0]['partitions'][1]['metadata'] is None
def test_parse_set_offset_response(self):
response = offset_commit_response()
val = self.client._parse_set_offset_response(response)
assert val == {'topic1': [0, 16]}
class TopicsTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.metadata_response = topic_metadata()
def test_maybe_update_full_metadata_expired(self):
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = 'metadata_response'
self.client._update_from_metadata = MagicMock()
fake_last_time = time.time() - (self.client.configuration.metadata_refresh * 2)
self.client._last_full_metadata = fake_last_time
self.client._maybe_update_full_metadata()
assert self.client._last_full_metadata > fake_last_time
self.client._send_any_broker.assert_called_once()
arg = self.client._send_any_broker.call_args[0][0]
assert isinstance(arg, TopicMetadataV1Request)
assert arg['topics'] is None
self.client._update_from_metadata.assert_called_once_with('metadata_response', delete=True)
def test_maybe_update_full_metadata_nocache(self):
self.client._send_any_broker = MagicMock()
self.client._update_from_metadata = MagicMock()
fake_last_time = time.time() - 1
self.client._last_full_metadata = fake_last_time
self.client._maybe_update_full_metadata(cache=False)
assert self.client._last_full_metadata > fake_last_time
def test_maybe_update_full_metadata_usecache(self):
self.client._send_any_broker = MagicMock()
self.client._update_from_metadata = MagicMock()
fake_last_time = time.time() - 1
self.client._last_full_metadata = fake_last_time
self.client._maybe_update_full_metadata(cache=True)
assert self.client._last_full_metadata == fake_last_time
def test_maybe_update_metadata_for_topics_noupdate(self):
self.client._update_from_metadata(self.metadata_response)
self.client.cluster.topics['topic1']._last_updated = time.time()
self.client._update_from_metadata = MagicMock()
self.client._send_any_broker = MagicMock()
self.client._maybe_update_metadata_for_topics(['topic1'])
self.client._update_from_metadata.assert_not_called()
def test_maybe_update_metadata_for_topics_expired(self):
self.client._update_from_metadata(self.metadata_response)
self.client.cluster.topics['topic1']._last_updated = 100
self.client._update_from_metadata = MagicMock()
self.client._send_any_broker = MagicMock()
self.client._maybe_update_metadata_for_topics(['topic1'])
self.client._update_from_metadata.assert_called_once()
req = self.client._send_any_broker.call_args[0][0]
assert len(req['topics']) == 1
assert req['topics'][0] == 'topic1'
def test_maybe_update_metadata_for_topics_forced(self):
self.client._update_from_metadata(self.metadata_response)
self.client.cluster.topics['topic1']._last_updated = time.time()
self.client._update_from_metadata = MagicMock()
self.client._send_any_broker = MagicMock()
self.client._maybe_update_metadata_for_topics(['topic1'], cache=False)
self.client._update_from_metadata.assert_called_once()
req = self.client._send_any_broker.call_args[0][0]
assert len(req['topics']) == 1
assert req['topics'][0] == 'topic1'
def test_maybe_update_metadata_for_topics_nonexistent(self):
self.client._update_from_metadata(self.metadata_response)
self.client.cluster.topics['topic1']._last_updated = time.time()
self.client._update_from_metadata = MagicMock()
self.client._send_any_broker = MagicMock()
self.client._maybe_update_metadata_for_topics(['topic1', 'topic2'])
self.client._update_from_metadata.assert_called_once()
req = self.client._send_any_broker.call_args[0][0]
assert len(req['topics']) == 2
assert req['topics'][0] == 'topic1'
assert req['topics'][1] == 'topic2'
def test_update_from_metadata(self):
self.client._update_brokers_from_metadata = MagicMock()
self.client._update_topics_from_metadata = MagicMock()
self.client._update_from_metadata('fake_metadata')
self.client._update_brokers_from_metadata.assert_called_once_with('fake_metadata')
self.client._update_topics_from_metadata.assert_called_once_with('fake_metadata', delete=False)
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 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_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_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_delete_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(broker2)
topic.partitions[0].add_replica(broker1)
topic.partitions[0].add_replica(broker3)
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_brokers_from_metadata(self):
self.client._update_brokers_from_metadata(self.metadata_response)
assert_cluster_has_brokers(self.client.cluster, self.metadata_response)
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_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 test_map_topic_partitions_to_brokers(self):
self.client._update_from_metadata(self.metadata_response)
val = self.client._map_topic_partitions_to_brokers(['topic1'])
assert val == {1: {'topic1': [0]}, 101: {'topic1': [1]}}
def test_map_topic_partitions_to_brokers_nonexistent(self):
self.client._update_from_metadata(self.metadata_response)
self.assertRaises(TopicError, self.client._map_topic_partitions_to_brokers, ['nosuchtopic'])
class InterfaceOffsetsTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.client._connected = True
# Get the broker and topic from a metadata update
self.client._update_from_metadata(topic_metadata())
self.group = Group('testgroup')
self.offset_fetch = offset_fetch()
def test_get_offsets_for_topics(self):
self.client._maybe_update_metadata_for_topics = MagicMock()
self.client._send_list_offsets_to_brokers = MagicMock()
self.client._send_list_offsets_to_brokers.return_value = {'topic1': 'responseobj'}
val = self.client.get_offsets_for_topics(['topic1'])
self.client._maybe_update_metadata_for_topics.assert_called_once_with(['topic1'])
self.client._send_list_offsets_to_brokers.assert_called_once()
values = self.client._send_list_offsets_to_brokers.call_args[0][0]
assert len(values) == 2
for broker_id in (1, 101):
assert values[broker_id]['replica_id'] == -1
assert len(values[broker_id]['topics']) == 1
assert values[broker_id]['topics'][0]['topic'] == 'topic1'
assert len(values[broker_id]['topics'][0]['partitions']) == 1
assert values[broker_id]['topics'][0]['partitions'][0]['timestamp'] == Client.OFFSET_LATEST
assert values[1]['topics'][0]['partitions'][0]['partition'] == 0
assert values[101]['topics'][0]['partitions'][0]['partition'] == 1
assert val == {'topic1': 'responseobj'}
def test_get_offsets_for_topic(self):
self.client.get_offsets_for_topics = MagicMock()
self.client.get_offsets_for_topics.return_value = {'topic1': 'thingitreturns'}
val = self.client.get_offsets_for_topic('topic1')
self.client.get_offsets_for_topics.assert_called_once_with(['topic1'], Client.OFFSET_LATEST)
assert val == 'thingitreturns'
def test_get_offsets_for_topic_bad_timestamp(self):
self.assertRaises(TypeError, self.client.get_offsets_for_topics, ['topic1'], timestamp='notanint')
def test_get_topics_for_group_string(self):
val = self.client._get_topics_for_group(self.group, 'testtopic')
assert val == ['testtopic']
def test_get_topics_for_group_list(self):
val = self.client._get_topics_for_group(self.group, ['topica', 'topicb'])
assert val == ['topica', 'topicb']
def test_get_topics_for_group_empty(self):
self.assertRaises(GroupError, self.client._get_topics_for_group, self.group, [])
def test_get_topics_for_group_subscribed(self):
self.group.subscribed_topics = MagicMock()
self.group.subscribed_topics.return_value = ['topica']
val = self.client._get_topics_for_group(self.group, None)
self.group.subscribed_topics.assert_called_once_with()
assert val == ['topica']
def test_get_offsets_for_group(self):
self.client.get_group = MagicMock()
self.client.get_group.return_value = self.group
self.client._get_topics_for_group = MagicMock()
self.client._get_topics_for_group.return_value = ['topic1']
self.client._maybe_update_metadata_for_topics = MagicMock()
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.return_value = self.offset_fetch
val = self.client.get_offsets_for_group('testgroup')
self.client.get_group.assert_called_once_with('testgroup')
self.client._get_topics_for_group.assert_called_once_with(self.group, None)
self.client._maybe_update_metadata_for_topics.assert_called_once_with(['topic1'])
self.client._send_group_aware_request.assert_called_once()
assert self.client._send_group_aware_request.call_args[0][0] == 'testgroup'
req = self.client._send_group_aware_request.call_args[0][1]
assert isinstance(req, OffsetFetchV1Request)
assert req['group_id'] == 'testgroup'
assert len(req['topics']) == 1
assert req['topics'][0]['topic'] == 'topic1'
assert len(req['topics'][0]['partitions']) == 2
assert req['topics'][0]['partitions'][0] == 0
assert req['topics'][0]['partitions'][1] == 1
assert 'topic1' in val
assert isinstance(val['topic1'], TopicOffsets)
assert val['topic1'].partitions == [4829, 8904]
def test_set_offsets_for_group_bad_offsets(self):
self.assertRaises(TypeError, self.client.set_offsets_for_group, 'testgroup', 'notalist')
def test_set_offsets_for_group_offsets_none(self):
self.assertRaises(TypeError, self.client.set_offsets_for_group, 'testgroup', None)
def test_set_offsets_for_group_not_empty(self):
self.group.state = 'Stable'
self.client.get_group = MagicMock()
self.client.get_group.return_value = self.group
self.assertRaises(GroupError, self.client.set_offsets_for_group, 'testgroup', [])
self.client.get_group.assert_called_once_with('testgroup')
def test_set_offsets_for_group(self):
self.client.get_group = MagicMock()
self.client.get_group.return_value = self.group
self.client._send_set_offset_request = MagicMock()
self.client._send_set_offset_request.return_value = 'sendresponse'
self.client._parse_set_offset_response = MagicMock()
self.client._parse_set_offset_response.return_value = {'topic1': [0, 0]}
offsets = TopicOffsets(self.client.cluster.topics['topic1'])
offsets.partitions[0] = 2342
offsets.partitions[1] = 8793
val = self.client.set_offsets_for_group('testgroup', [offsets])
assert val == {'topic1': [0, 0]}
self.client.get_group.assert_called_once_with('testgroup')
self.client._send_set_offset_request.assert_called_once_with('testgroup', [offsets])
self.client._parse_set_offset_response.assert_called_once_with('sendresponse')
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client(
broker_list='broker1.example.com:9091,broker2.example.com:9092')
self.client._connected = True
self.metadata_response = topic_metadata()
class SendHelperTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.group_coordinator = group_coordinator()
self.coordinator_error = group_coordinator_error()
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')
@patch('kafka.tools.client.Broker')
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_send_group_aware_request_error(self):
self.client._send_any_broker = MagicMock()
self.client._send_any_broker.return_value = self.coordinator_error
self.assertRaises(GroupError, self.client._send_group_aware_request,
'testgroup', 'fakerequest')
def test_send_all_brokers(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_all_brokers('fakerequest')
broker1.send.assert_called_once_with('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
assert val[1] == 'fakeresponse'
assert val[101] == 'otherresponse'
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
@patch('kafka.tools.client.shuffle', lambda x: sorted(x))
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'
@patch('kafka.tools.client.shuffle', lambda x: sorted(x))
def test_send_any_broker_one_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_any_broker('fakerequest')
broker2.send.assert_called_once_with('fakerequest')
broker1.send.assert_called_once_with('fakerequest')
assert val == 'fakeresponse'
@patch('kafka.tools.client.shuffle', lambda x: sorted(x))
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_raise_if_not_connected(self):
self.assertRaises(ConnectionError, self.client._raise_if_not_connected)
def test_connect_zookeeper(self, mock_create):
mock_create.return_value = Cluster()
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
test_client.connect()
mock_create.assert_called_once_with(zkconnect='zk.example.com:2181/kafka-cluster', fetch_topics=False)
class OffsetTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
# Get the broker and topic from a metadata update
self.client._update_from_metadata(topic_metadata())
self.list_offset = list_offset()
self.list_offset_error = list_offset_error()
self.list_offset_request = {
1: {
'replica_id':
-1,
'topics': [{
'topic':
'topic1',
'partitions': [{
'partition': 0,
'timestamp': Client.OFFSET_LATEST,
'max_num_offsets': 1
}, {
'partition': 1,
'timestamp': Client.OFFSET_EARLIEST,
'max_num_offsets': 1
}]
}]
}
}
def test_send_list_offsets_to_brokers(self):
self.client.cluster.brokers[1].send = MagicMock()
self.client.cluster.brokers[1].send.return_value = (1,
self.list_offset)
val = self.client._send_list_offsets_to_brokers(
self.list_offset_request)
self.client.cluster.brokers[1].send.assert_called_once()
req = self.client.cluster.brokers[1].send.call_args[0][0]
assert req['replica_id'] == -1
assert len(req['topics']) == 1
assert req['topics'][0]['topic'] == 'topic1'
assert len(req['topics'][0]['partitions']) == 2
assert req['topics'][0]['partitions'][0]['partition'] == 0
assert req['topics'][0]['partitions'][0][
'timestamp'] == Client.OFFSET_LATEST
assert req['topics'][0]['partitions'][0]['max_num_offsets'] == 1
assert req['topics'][0]['partitions'][1]['partition'] == 1
assert req['topics'][0]['partitions'][1][
'timestamp'] == Client.OFFSET_EARLIEST
assert req['topics'][0]['partitions'][1]['max_num_offsets'] == 1
assert len(val) == 1
assert 'topic1' in val
assert isinstance(val['topic1'], TopicOffsets)
assert val['topic1'].topic == self.client.cluster.topics['topic1']
assert len(val['topic1'].partitions) == 2
assert val['topic1'].partitions[0] == 4829
assert val['topic1'].partitions[1] == 8904
def test_send_list_offsets_to_brokers_connection_error(self):
self.client.cluster.brokers[1].send = MagicMock()
self.client.cluster.brokers[1].send.side_effect = ConnectionError
self.assertRaises(ConnectionError,
self.client._send_list_offsets_to_brokers,
self.list_offset_request)
def test_send_list_offsets_to_brokers_offsets_error(self):
self.client.cluster.brokers[1].send = MagicMock()
self.client.cluster.brokers[1].send.return_value = (
1, self.list_offset_error)
self.assertRaises(OffsetError,
self.client._send_list_offsets_to_brokers,
self.list_offset_request)
def test_send_set_offset_request_bad_offsets(self):
self.assertRaises(TypeError, self.client._send_set_offset_request,
'testgroup', ['notatopicoffsets'])
def test_send_set_offset_request_bad_topic(self):
offsets = TopicOffsets(self.client.cluster.topics['topic1'])
offsets.topic = None
self.assertRaises(TypeError, self.client._send_set_offset_request,
'testgroup', [offsets])
def test_send_set_offset_request(self):
offsets = TopicOffsets(self.client.cluster.topics['topic1'])
offsets.partitions[0] = 2342
offsets.partitions[1] = 8793
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.return_value = 'responseobject'
val = self.client._send_set_offset_request('testgroup', [offsets])
assert val == 'responseobject'
self.client._send_group_aware_request.assert_called_once()
assert self.client._send_group_aware_request.call_args[0][
0] == 'testgroup'
req = self.client._send_group_aware_request.call_args[0][1]
assert isinstance(req, OffsetCommitV2Request)
assert req['group_id'] == 'testgroup'
assert req['group_generation_id'] == -1
assert req['member_id'] == ''
assert req['retention_time'] == -1
assert len(req['topics']) == 1
assert req['topics'][0]['topic'] == 'topic1'
assert len(req['topics'][0]['partitions']) == 2
assert req['topics'][0]['partitions'][0]['partition'] == 0
assert req['topics'][0]['partitions'][0]['offset'] == 2342
assert req['topics'][0]['partitions'][0]['metadata'] is None
assert req['topics'][0]['partitions'][1]['partition'] == 1
assert req['topics'][0]['partitions'][1]['offset'] == 8793
assert req['topics'][0]['partitions'][1]['metadata'] is None
def test_parse_set_offset_response(self):
response = offset_commit_response()
val = self.client._parse_set_offset_response(response)
assert val == {'topic1': [0, 16]}
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client(broker_list='broker1.example.com:9091,broker2.example.com:9092')
self.client._connected = True
self.metadata_response = topic_metadata()
class GenericInterfaceTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client(broker_list='broker1.example.com:9091,broker2.example.com:9092')
self.client._connected = True
self.metadata_response = topic_metadata()
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_create(self):
assert self.client.configuration.broker_list == [('broker1.example.com', 9091), ('broker2.example.com', 9092)]
def test_create_with_configuration(self):
config = ClientConfiguration(zkconnect='zk.example.com:2181/kafka-cluster')
client = Client(configuration=config)
assert client.configuration == config
def test_create_config_zkconnect(self):
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
assert test_client.configuration.zkconnect == 'zk.example.com:2181/kafka-cluster'
def test_create_config_bad_arg(self):
self.assertRaises(ConfigurationError, Client, invalidconfig='foo')
def test_create_config_bad_config_object(self):
self.assertRaises(ConfigurationError, Client, configuration='foo')
@patch.object(Cluster, 'create_from_zookeeper')
def test_connect_zookeeper(self, mock_create):
mock_create.return_value = Cluster()
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
test_client.connect()
mock_create.assert_called_once_with(zkconnect='zk.example.com:2181/kafka-cluster', fetch_topics=False)
@patch.object(Broker, 'connect')
@patch.object(Broker, 'send')
@patch.object(Broker, 'close')
def test_connect_broker_list(self, mock_close, mock_send, mock_connect):
def add_brokers_with_mocks(metadata):
self.client._update_brokers_from_metadata(metadata)
for broker_id in self.client.cluster.brokers:
self.client.cluster.brokers[broker_id].connect = MagicMock()
mock_connect.side_effect = [ConnectionError, None]
mock_send.return_value = (1, self.metadata_response)
self.client._update_from_metadata = MagicMock()
self.client._update_from_metadata.side_effect = add_brokers_with_mocks
self.client.connect()
mock_connect.assert_has_calls([call(), call()])
mock_send.assert_called_once()
assert isinstance(mock_send.call_args[0][0], TopicMetadataV1Request)
mock_close.assert_called_once()
# Metadata is updated and brokers are connected
self.client._update_from_metadata.assert_called_once_with(self.metadata_response)
for broker_id in self.client.cluster.brokers:
self.client.cluster.brokers[broker_id].connect.assert_called_once()
def test_connect_broker_list_exhausted(self):
self.client._maybe_bootstrap_cluster = MagicMock()
self.client._maybe_bootstrap_cluster.return_value = False
self.assertRaises(ConnectionError, self.client.connect)
def test_connect_zookeeper(self, mock_create):
mock_create.return_value = Cluster()
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
test_client.connect()
mock_create.assert_called_once_with(
zkconnect='zk.example.com:2181/kafka-cluster', fetch_topics=False)
class GroupsTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.describe_groups = describe_groups()
def test_update_groups_from_describe_create(self):
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
def test_update_groups_from_describe_assignment(self):
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
group = self.client.cluster.groups['testgroup']
assert group.members[0].topics == {'topic1': [0]}
assert group.members[1].topics == {'topic1': [1]}
def test_update_groups_from_describe_update(self):
self.client.cluster.add_group(Group('testgroup'))
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
def test_update_groups_from_describe_clear_members(self):
self.client.cluster.add_group(Group('testgroup'))
self.client.cluster.groups['testgroup'].members = [GroupMember('badmember')]
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
def test_maybe_update_groups_list_expired(self):
self.client._send_all_brokers = MagicMock()
self.client._send_all_brokers.return_value = ['metadata_response']
self.client._update_groups_from_lists = MagicMock()
fake_last_time = time.time() - (self.client.configuration.metadata_refresh * 2)
self.client._last_group_list = fake_last_time
self.client._maybe_update_groups_list()
assert self.client._last_group_list > fake_last_time
self.client._send_all_brokers.assert_called_once()
arg = self.client._send_all_brokers.call_args[0][0]
assert isinstance(arg, ListGroupsV0Request)
self.client._update_groups_from_lists.assert_called_once_with(['metadata_response'])
def test_maybe_update_groups_list_nocache(self):
self.client._send_all_brokers = MagicMock()
self.client._update_groups_from_lists = MagicMock()
fake_last_time = time.time() - 1000
self.client._last_group_list = fake_last_time
self.client._maybe_update_groups_list(cache=False)
assert self.client._last_group_list > fake_last_time
def test_maybe_update_groups_list_usecache(self):
self.client._send_all_brokers = MagicMock()
self.client._update_groups_from_lists = MagicMock()
fake_last_time = time.time() - 1000
self.client._last_group_list = fake_last_time
self.client._maybe_update_groups_list(cache=True)
def test_update_groups_from_lists(self):
self.client._add_or_update_group = MagicMock()
list_group = list_groups()
val = self.client._update_groups_from_lists({1: list_group})
assert val == 0
self.client._add_or_update_group.assert_called_once_with(list_group['groups'][0], 1)
def test_update_groups_from_lists_error(self):
self.client._add_or_update_group = MagicMock()
val = self.client._update_groups_from_lists({1: list_groups_error()})
assert val == 1
self.client._add_or_update_group.assert_not_called()
def test_update_groups_from_lists_none(self):
self.client._add_or_update_group = MagicMock()
val = self.client._update_groups_from_lists({1: None})
assert val == 1
self.client._add_or_update_group.assert_not_called()
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_add_or_update_group_update(self):
broker = Broker('host1.example.com', id=1, port=8031)
self.client.cluster.add_broker(broker)
group = Group('group1')
self.client.cluster.add_group(group)
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'
class InterfaceTopicsTests(unittest.TestCase):
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_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_list_topics_cache(self):
self.client._maybe_update_full_metadata = MagicMock()
self.client.list_topics(cache=False)
self.client._maybe_update_full_metadata.assert_called_once_with(False)
def test_get_topic(self):
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_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 test_get_topic_force_cache(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
val = 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)
assert isinstance(val, Topic)
assert val.name == 'topic1'
assert len(val.partitions) == 2
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_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 setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.describe_groups = describe_groups()
class GenericInterfaceTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client(
broker_list='broker1.example.com:9091,broker2.example.com:9092')
self.client._connected = True
self.metadata_response = topic_metadata()
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_create(self):
assert self.client.configuration.broker_list == [
('broker1.example.com', 9091), ('broker2.example.com', 9092)
]
def test_create_with_configuration(self):
config = ClientConfiguration(
zkconnect='zk.example.com:2181/kafka-cluster')
client = Client(configuration=config)
assert client.configuration == config
def test_create_config_zkconnect(self):
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
assert test_client.configuration.zkconnect == 'zk.example.com:2181/kafka-cluster'
def test_create_config_bad_arg(self):
self.assertRaises(ConfigurationError, Client, invalidconfig='foo')
def test_create_config_bad_config_object(self):
self.assertRaises(ConfigurationError, Client, configuration='foo')
@patch.object(Cluster, 'create_from_zookeeper')
def test_connect_zookeeper(self, mock_create):
mock_create.return_value = Cluster()
test_client = Client(zkconnect='zk.example.com:2181/kafka-cluster')
test_client.connect()
mock_create.assert_called_once_with(
zkconnect='zk.example.com:2181/kafka-cluster', fetch_topics=False)
@patch.object(Broker, 'connect')
@patch.object(Broker, 'send')
@patch.object(Broker, 'close')
def test_connect_broker_list(self, mock_close, mock_send, mock_connect):
def add_brokers_with_mocks(metadata):
self.client._update_brokers_from_metadata(metadata)
for broker_id in self.client.cluster.brokers:
self.client.cluster.brokers[broker_id].connect = MagicMock()
mock_connect.side_effect = [ConnectionError, None]
mock_send.return_value = (1, self.metadata_response)
self.client._update_from_metadata = MagicMock()
self.client._update_from_metadata.side_effect = add_brokers_with_mocks
self.client.connect()
mock_connect.assert_has_calls([call(), call()])
mock_send.assert_called_once()
assert isinstance(mock_send.call_args[0][0], TopicMetadataV1Request)
mock_close.assert_called_once()
# Metadata is updated and brokers are connected
self.client._update_from_metadata.assert_called_once_with(
self.metadata_response)
for broker_id in self.client.cluster.brokers:
self.client.cluster.brokers[broker_id].connect.assert_called_once()
def test_connect_broker_list_exhausted(self):
self.client._maybe_bootstrap_cluster = MagicMock()
self.client._maybe_bootstrap_cluster.return_value = False
self.assertRaises(ConnectionError, self.client.connect)
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.metadata_response = topic_metadata()
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.describe_groups = describe_groups()
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.group_coordinator = group_coordinator()
self.coordinator_error = group_coordinator_error()
def test_create_with_configuration(self):
config = ClientConfiguration(
zkconnect='zk.example.com:2181/kafka-cluster')
client = Client(configuration=config)
assert client.configuration == config
class GroupsTests(unittest.TestCase):
def setUp(self):
# Dummy client for testing - we're not going to connect that bootstrap broker
self.client = Client()
self.describe_groups = describe_groups()
def test_update_groups_from_describe_create(self):
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
def test_update_groups_from_describe_assignment(self):
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
group = self.client.cluster.groups['testgroup']
assert group.members[0].topics == {'topic1': [0]}
assert group.members[1].topics == {'topic1': [1]}
def test_update_groups_from_describe_update(self):
self.client.cluster.add_group(Group('testgroup'))
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
def test_update_groups_from_describe_clear_members(self):
self.client.cluster.add_group(Group('testgroup'))
self.client.cluster.groups['testgroup'].members = [GroupMember('badmember')]
self.client._update_groups_from_describe(self.describe_groups)
assert_cluster_has_groups(self.client.cluster, self.describe_groups)
def test_maybe_update_groups_list_expired(self):
self.client._send_all_brokers = MagicMock()
self.client._send_all_brokers.return_value = ['metadata_response']
self.client._update_groups_from_lists = MagicMock()
fake_last_time = time.time() - (self.client.configuration.metadata_refresh * 2)
self.client._last_group_list = fake_last_time
self.client._maybe_update_groups_list()
assert self.client._last_group_list > fake_last_time
self.client._send_all_brokers.assert_called_once()
arg = self.client._send_all_brokers.call_args[0][0]
assert isinstance(arg, ListGroupsV0Request)
self.client._update_groups_from_lists.assert_called_once_with(['metadata_response'])
def test_maybe_update_groups_list_nocache(self):
self.client._send_all_brokers = MagicMock()
self.client._update_groups_from_lists = MagicMock()
fake_last_time = time.time() - 1
self.client._last_group_list = fake_last_time
self.client._maybe_update_groups_list(cache=False)
assert self.client._last_group_list > fake_last_time
def test_maybe_update_groups_list_usecache(self):
self.client._send_all_brokers = MagicMock()
self.client._update_groups_from_lists = MagicMock()
fake_last_time = time.time() - 1
self.client._last_group_list = fake_last_time
self.client._maybe_update_groups_list(cache=True)
def test_update_groups_from_lists(self):
self.client._add_or_update_group = MagicMock()
list_group = list_groups()
val = self.client._update_groups_from_lists({1: list_group})
assert val == 0
self.client._add_or_update_group.assert_called_once_with(list_group['groups'][0], 1)
def test_update_groups_from_lists_error(self):
self.client._add_or_update_group = MagicMock()
val = self.client._update_groups_from_lists({1: list_groups_error()})
assert val == 1
self.client._add_or_update_group.assert_not_called()
def test_update_groups_from_lists_none(self):
self.client._add_or_update_group = MagicMock()
val = self.client._update_groups_from_lists({1: None})
assert val == 1
self.client._add_or_update_group.assert_not_called()
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_add_or_update_group_update(self):
broker = Broker('host1.example.com', id=1, port=8031)
self.client.cluster.add_broker(broker)
group = Group('group1')
self.client.cluster.add_group(group)
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'
class InterfaceGroupsTests(unittest.TestCase):
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.describe_groups = describe_groups()
self.describe_groups_error = describe_groups_error()
def test_list_groups(self):
self.client.cluster.add_group(Group('group2'))
self.client._maybe_update_groups_list = MagicMock()
self.client._maybe_update_groups_list.return_value = 312
groups, errs = self.client.list_groups()
self.client._maybe_update_groups_list.assert_called_once_with(True)
assert errs == 312
assert groups == ['group2']
def test_list_groups_cache(self):
self.client._maybe_update_groups_list = MagicMock()
self.client.list_groups(cache=False)
self.client._maybe_update_groups_list.assert_called_once_with(False)
def test_get_group(self):
def add_group_to_cluster(group_name, request):
# _send_group_aware_request has a side effect of creating the group and setting the coordinator
self.client.cluster.add_group(Group(group_name))
self.client.cluster.groups[group_name].coordinator = self.client.cluster.brokers[1]
return self.describe_groups
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.side_effect = add_group_to_cluster
val = self.client.get_group('testgroup')
self.client._send_group_aware_request.assert_called_once()
assert self.client._send_group_aware_request.call_args[0][0] == 'testgroup'
assert isinstance(self.client._send_group_aware_request.call_args[0][1], DescribeGroupsV0Request)
assert isinstance(val, Group)
assert val.name == 'testgroup'
print(val.coordinator)
print(self.client.cluster.brokers[1])
assert val.coordinator == self.client.cluster.brokers[1]
def test_get_group_existing_cached(self):
group = Group('testgroup')
group.coordinator = self.client.cluster.brokers[1]
group._last_updated = time.time()
self.client.cluster.add_group(group)
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.return_value = self.describe_groups
self.client._update_groups_from_describe = MagicMock()
val = self.client.get_group('testgroup')
self.client._send_group_aware_request.assert_not_called()
self.client._update_groups_from_describe.assert_not_called()
assert isinstance(val, Group)
assert val.name == 'testgroup'
assert val.coordinator == self.client.cluster.brokers[1]
def test_get_group_error(self):
group = Group('badgroup')
group.coordinator = self.client.cluster.brokers[1]
self.client.cluster.add_group(group)
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.return_value = self.describe_groups_error
self.assertRaises(GroupError, self.client.get_group, 'badgroup', cache=False)
self.client._send_group_aware_request.assert_called_once()
def test_get_group_nonexistent(self):
self.client._send_group_aware_request = MagicMock()
self.client._send_group_aware_request.side_effect = GroupError
self.assertRaises(GroupError, self.client.get_group, 'nonexistentgroup')
self.client._send_group_aware_request.assert_called_once()