def test_process_error(self, config):
message_iterator = iter([
Message(1, 12345, 'key1', 'value1'),
Message(1, 12346, 'key2', 'value2'),
Message(1, 12347, 'key1', 'value3'),
])
with mock_kafka() as (mock_client, _):
with mock.patch.object(
KafkaSimpleConsumer,
'__iter__',
return_value=message_iterator,
):
consumer = KafkaConsumerBase('test_topic', config)
consumer.process = mock.Mock(side_effect=Exception('Boom!'))
consumer.initialize = mock.Mock()
consumer.dispose = mock.Mock()
with pytest.raises(ProcessMessageError):
consumer.run()
def test_commit_message_error(self, config):
with mock_kafka() as (mock_client, mock_consumer):
consumer = KafkaSimpleConsumer('test_topic', config)
consumer.connect()
mock_client.return_value.send_offset_commit_request \
.side_effect = KafkaError("Boom!")
actual = consumer.commit_message(
Message(0, 100, 'mykey', 'myvalue'), )
assert actual is False
def test_run_and_terminate(self, config):
message_iterator = iter([
Message(1, 12345, 'key1', 'value1'),
Message(1, 12346, 'key2', 'value2'),
Message(1, 12347, 'key1', 'value3'),
])
with mock_kafka() as (mock_client, mock_consumer):
with mock.patch.object(KafkaSimpleConsumer,
'__iter__',
return_value=message_iterator):
with mock.patch.object(KafkaSimpleConsumer,
'commit') as mock_commit:
consumer = KafkaConsumerBase('test_topic', config)
consumer.process = mock.Mock()
consumer.initialize = mock.Mock()
consumer.dispose = mock.Mock()
consumer.terminate()
consumer.run()
assert consumer.initialize.call_count == 1
# process should have been called 0 times
# termination flag is checked before the first
# message is pulled.
assert consumer.process.call_count == 0
# check just last call arguments
consumer.process.calls_args_list([
Message(1, 12347, 'key1', 'value3'),
Message(1, 12345, 'key1', 'value1'),
Message(1, 12346, 'key2', 'value2'),
])
consumer.dispose.assert_called_once_with()
mock_commit.assert_called_once_with()
mock_client.return_value.close.assert_called_once_with()
def test_commit_message_default(self, config):
with mock_kafka() as (mock_client, mock_consumer):
consumer = KafkaSimpleConsumer('test_topic', config)
consumer.connect()
actual = consumer.commit_message(
Message(0, 100, 'mykey', 'myvalue'), )
assert actual is True
mock_client.return_value.send_offset_commit_request \
.assert_called_once_with(
'test_group'.encode(),
[OffsetCommitRequest('test_topic'.encode(), 0, 100, None)],
)
def test_commit_message_zk(self, config):
if getattr(KafkaClient, 'send_offset_commit_request_kafka',
None) is None:
return
with mock_kafka() as (mock_client, mock_consumer):
config._config['offset_storage'] = 'zookeeper'
consumer = KafkaSimpleConsumer('test_topic', config)
consumer.connect()
actual = consumer.commit_message(
Message(0, 100, 'mykey', 'myvalue'), )
assert actual is True
mock_client.return_value.send_offset_commit_request \
.assert_called_once_with(
'test_group'.encode(),
[OffsetCommitRequest('test_topic'.encode(), 0, 100, None)],
)
def test_get_message(self, config):
with mock_kafka() as (_, mock_consumer):
mock_obj = mock_consumer.return_value
# get message should return a tuple (partition_id, (offset,
# Message)). Message is a namedtuple defined in
# kafka-python that at least contains key and value.
mock_message = mock.Mock()
mock_message.value = 'test_content'
mock_message.key = 'test_key'
kafka_message = (1, (12345, mock_message))
mock_obj.get_message.return_value = kafka_message
# Set the current offset the offset of the message + 1
mock_obj.offsets = {1: 12346}
consumer = KafkaSimpleConsumer('test_topic', config)
consumer.connect()
assert consumer.get_message() == Message(
partition=1,
offset=12345,
key='test_key',
value='test_content',
)