def test__release_retry(self, cluster):
config = KafkaConsumerConfig(self.group,
cluster,
auto_commit_enable=True)
consumer = KafkaConsumerGroup([], config)
mock_consumer = mock.Mock()
mock_consumer.set_topic_partitions.side_effect = KafkaUnavailableError
consumer.consumer = mock_consumer
with pytest.raises(KafkaUnavailableError):
consumer._release({})
assert mock_consumer.set_topic_partitions.call_count == 2
def test__acquire_has_consumer(self, cluster, example_partitions,
mock_post_rebalance_cb):
config = KafkaConsumerConfig(
self.group,
cluster,
post_rebalance_callback=mock_post_rebalance_cb)
consumer = KafkaConsumerGroup([], config)
consumer.consumer = mock.Mock()
consumer._acquire(example_partitions)
consumer.consumer.set_topic_partitions.assert_called_once_with(
example_partitions)
mock_post_rebalance_cb.assert_called_once_with(example_partitions)
def test__release(self, cluster, example_partitions,
mock_pre_rebalance_cb):
config = KafkaConsumerConfig(
self.group,
cluster,
auto_commit_enable=True,
pre_rebalance_callback=mock_pre_rebalance_cb)
consumer = KafkaConsumerGroup([], config)
mock_consumer = mock.Mock()
consumer.consumer = mock_consumer
consumer._release(example_partitions)
mock_consumer.commit.assert_called_once_with()
mock_consumer.set_topic_partitions.assert_called_once_with({})
mock_pre_rebalance_cb.assert_called_once_with(example_partitions)
def test_next(self, mock_consumer, mock_partitioner, cluster):
config = KafkaConsumerConfig(self.group,
cluster,
consumer_timeout_ms=500)
consumer = KafkaConsumerGroup([], config)
consumer.partitioner = mock_partitioner()
consumer.consumer = mock_consumer()
def fake_next():
time.sleep(1)
raise ConsumerTimeout()
consumer.consumer.next.side_effect = fake_next
# The mock KafkaConsumer.next (called fake_next above) takes longer than
# consumer_timeout_ms, so we should get a ConsumerTimeout from
# KafkaConsumerGroup
with pytest.raises(ConsumerTimeout):
consumer.next()
consumer.consumer.next.assert_called_once_with()
consumer.partitioner.refresh.assert_called_once_with()