def test__should_keep_trying_no_timeout(self, cluster):
config = KafkaConsumerConfig(self.group,
cluster,
consumer_timeout_ms=-1)
consumer = KafkaConsumerGroup([], config)
long_time_ago = time.time() - 1000
assert consumer._should_keep_trying(long_time_ago)
def test__acquire_has_no_consumer(self, mock_consumer, cluster,
example_partitions):
config = KafkaConsumerConfig(self.group, cluster)
consumer = KafkaConsumerGroup([], config)
consumer._acquire(example_partitions)
mock_consumer.assert_called_once_with(example_partitions,
**consumer.config)
def consume():
consumer = KafkaConsumerGroup([topic], config)
with consumer:
while True:
try:
message = consumer.next()
queue.put(message)
consumer.task_done(message)
except ConsumerTimeout:
return
def test__should_keep_trying_timed_out(self, mock_time, cluster):
mock_time.return_value = 0
config = KafkaConsumerConfig(self.group,
cluster,
consumer_timeout_ms=1000)
consumer = KafkaConsumerGroup([], config)
over_a_second_ago = time.time() - 1.2
assert not consumer._should_keep_trying(over_a_second_ago)
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()
def _start(self):
self.consumer_group = KafkaConsumerGroup(
topics=self.topic_to_partition_map.keys(),
config=self._kafka_consumer_config
)
self.consumer_group.start()
def test__auto_commit_enabled_not_enabled(self, cluster):
config = KafkaConsumerConfig(self.group,
cluster,
auto_commit_enable=False)
consumer = KafkaConsumerGroup([], config)
assert not consumer._auto_commit_enabled()
def test___init__string_topics(self):
with pytest.raises(AssertionError):
KafkaConsumerGroup(self.topic, None)