def test_unknown_topic_or_partition(self):
consumer = AIOKafkaConsumer(
loop=self.loop, group_id=None,
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=('some_topic_unknown',))
with self.assertRaises(UnknownTopicOrPartitionError):
yield from consumer.start()
with self.assertRaises(UnknownTopicOrPartitionError):
yield from consumer.assign([TopicPartition(self.topic, 2222)])
def test_unknown_topic_or_partition(self):
consumer = AIOKafkaConsumer(loop=self.loop,
group_id=None,
bootstrap_servers=self.hosts,
auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=('some_topic_unknown', ))
with self.assertRaises(UnknownTopicOrPartitionError):
yield from consumer.start()
with self.assertRaises(UnknownTopicOrPartitionError):
yield from consumer.assign([TopicPartition(self.topic, 2222)])
def test_manual_subscribe_nogroup(self):
msgs1 = yield from self.send_messages(0, range(0, 10))
msgs2 = yield from self.send_messages(1, range(10, 20))
available_msgs = msgs1 + msgs2
consumer = AIOKafkaConsumer(
loop=self.loop, group_id=None,
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=(self.topic,))
yield from consumer.start()
result = []
for i in range(20):
msg = yield from consumer.getone()
result.append(msg.value)
self.assertEqual(set(available_msgs), set(result))
yield from consumer.stop()
def test_manual_subscribe_nogroup(self):
msgs1 = yield from self.send_messages(0, range(0, 10))
msgs2 = yield from self.send_messages(1, range(10, 20))
available_msgs = msgs1 + msgs2
consumer = AIOKafkaConsumer(
loop=self.loop, group_id=None,
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=(self.topic,))
yield from consumer.start()
result = []
for i in range(20):
msg = yield from consumer.getone()
result.append(msg.value)
self.assertEqual(set(available_msgs), set(result))
yield from consumer.stop()
def test_manual_subscribe_pattern(self):
msgs1 = yield from self.send_messages(0, range(0, 10))
msgs2 = yield from self.send_messages(1, range(10, 20))
available_msgs = msgs1 + msgs2
consumer = AIOKafkaConsumer(loop=self.loop,
group_id='test-group',
bootstrap_servers=self.hosts,
auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(pattern="topic-test_manual_subs*")
yield from consumer.start()
yield from consumer.seek_to_committed()
result = []
for i in range(20):
msg = yield from consumer.getone()
result.append(msg.value)
self.assertEqual(set(available_msgs), set(result))
yield from consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
yield from consumer.seek_to_committed(TopicPartition(self.topic, 0))
msg = yield from consumer.getone(TopicPartition(self.topic, 0))
self.assertEqual(msg.value, b'9')
yield from consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
yield from consumer.stop()
# subscribe by topic
consumer = AIOKafkaConsumer(loop=self.loop,
group_id='test-group',
bootstrap_servers=self.hosts,
auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=(self.topic, ))
yield from consumer.start()
yield from consumer.seek_to_committed()
result = []
for i in range(10):
msg = yield from consumer.getone()
result.append(msg.value)
self.assertEqual(set(msgs2), set(result))
self.assertEqual(consumer.subscription(), set([self.topic]))
yield from consumer.stop()
def test_consumer_subscribe_pattern_autocreate_no_group_id(self):
pattern = "^no-group-pattern-.*$"
consumer = AIOKafkaConsumer(loop=self.loop,
bootstrap_servers=self.hosts,
metadata_max_age_ms=200,
group_id=None,
fetch_max_wait_ms=50,
auto_offset_reset="earliest")
self.add_cleanup(consumer.stop)
yield from consumer.start()
consumer.subscribe(pattern=pattern)
# Start getter for the topics. Should not create any topics
consume_task = self.loop.create_task(consumer.getone())
yield from asyncio.sleep(0.3, loop=self.loop)
self.assertFalse(consume_task.done())
self.assertEqual(consumer.subscription(), set())
# Now lets autocreate the topic by fetching metadata for it.
producer = AIOKafkaProducer(loop=self.loop,
bootstrap_servers=self.hosts)
self.add_cleanup(producer.stop)
yield from producer.start()
my_topic = "no-group-pattern-1"
yield from producer.client._wait_on_metadata(my_topic)
# Wait for consumer to refresh metadata with new topic
yield from asyncio.sleep(0.3, loop=self.loop)
self.assertFalse(consume_task.done())
self.assertTrue(consumer._client.cluster.topics() >= {my_topic})
self.assertEqual(consumer.subscription(), {my_topic})
# Add another topic
my_topic2 = "no-group-pattern-2"
yield from producer.client._wait_on_metadata(my_topic2)
# Wait for consumer to refresh metadata with new topic
yield from asyncio.sleep(0.3, loop=self.loop)
self.assertFalse(consume_task.done())
self.assertTrue(
consumer._client.cluster.topics() >= {my_topic, my_topic2})
self.assertEqual(consumer.subscription(), {my_topic, my_topic2})
# Now lets actualy produce some data and verify that it is consumed
yield from producer.send(my_topic, b'test msg')
data = yield from consume_task
self.assertEqual(data.value, b'test msg')
def test_manual_subscribe_pattern(self):
msgs1 = yield from self.send_messages(0, range(0, 10))
msgs2 = yield from self.send_messages(1, range(10, 20))
available_msgs = msgs1 + msgs2
consumer = AIOKafkaConsumer(
loop=self.loop, group_id='test-group',
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(pattern="topic-test_manual_subs*")
yield from consumer.start()
yield from consumer.seek_to_committed()
result = []
for i in range(20):
msg = yield from consumer.getone()
result.append(msg.value)
self.assertEqual(set(available_msgs), set(result))
yield from consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(9, '')})
yield from consumer.seek_to_committed(TopicPartition(self.topic, 0))
msg = yield from consumer.getone()
self.assertEqual(msg.value, b'9')
yield from consumer.commit(
{TopicPartition(self.topic, 0): OffsetAndMetadata(10, '')})
yield from consumer.stop()
# subscribe by topic
consumer = AIOKafkaConsumer(
loop=self.loop, group_id='test-group',
bootstrap_servers=self.hosts, auto_offset_reset='earliest',
enable_auto_commit=False)
consumer.subscribe(topics=(self.topic,))
yield from consumer.start()
yield from consumer.seek_to_committed()
result = []
for i in range(10):
msg = yield from consumer.getone()
result.append(msg.value)
self.assertEqual(set(msgs2), set(result))
self.assertEqual(consumer.subscription(), set([self.topic]))
yield from consumer.stop()
def test_exclude_internal_topics(self):
# Create random topic
my_topic = "some_noninternal_topic"
client = AIOKafkaClient(
loop=self.loop, bootstrap_servers=self.hosts,
client_id="test_autocreate")
yield from client.bootstrap()
yield from client._wait_on_metadata(my_topic)
yield from client.close()
# Check if only it will be subscribed
pattern = "^.*$"
consumer = AIOKafkaConsumer(
loop=self.loop, bootstrap_servers=self.hosts,
metadata_max_age_ms=200, group_id="some_group_1",
auto_offset_reset="earliest",
exclude_internal_topics=False)
consumer.subscribe(pattern=pattern)
yield from consumer.start()
self.assertIn("__consumer_offsets", consumer.subscription())
yield from consumer._client.force_metadata_update()
self.assertIn("__consumer_offsets", consumer.subscription())
yield from consumer.stop()
def test_consumer_rebalance_on_new_topic(self):
# Test will create a consumer group and check if adding new topic
# will trigger a group rebalance and assign partitions
pattern = "^another-autocreate-pattern-.*$"
client = AIOKafkaClient(
loop=self.loop, bootstrap_servers=self.hosts,
client_id="test_autocreate")
yield from client.bootstrap()
listener1 = StubRebalanceListener(loop=self.loop)
listener2 = StubRebalanceListener(loop=self.loop)
consumer1 = AIOKafkaConsumer(
loop=self.loop, bootstrap_servers=self.hosts,
metadata_max_age_ms=200, group_id="test-autocreate-rebalance",
heartbeat_interval_ms=100)
consumer1.subscribe(pattern=pattern, listener=listener1)
yield from consumer1.start()
consumer2 = AIOKafkaConsumer(
loop=self.loop, bootstrap_servers=self.hosts,
metadata_max_age_ms=200, group_id="test-autocreate-rebalance",
heartbeat_interval_ms=100)
consumer2.subscribe(pattern=pattern, listener=listener2)
yield from consumer2.start()
yield from asyncio.sleep(0.5, loop=self.loop)
# bootstrap will take care of the initial group assignment
self.assertEqual(consumer1.assignment(), set())
self.assertEqual(consumer2.assignment(), set())
listener1.reset()
listener2.reset()
# Lets force autocreation of a topic
my_topic = "another-autocreate-pattern-1"
yield from client._wait_on_metadata(my_topic)
# Wait for group to stabilize
assign1 = yield from listener1.wait_assign()
assign2 = yield from listener2.wait_assign()
# We expect 2 partitons for autocreated topics
my_partitions = set([
TopicPartition(my_topic, 0), TopicPartition(my_topic, 1)])
self.assertEqual(assign1 | assign2, my_partitions)
self.assertEqual(
consumer1.assignment() | consumer2.assignment(),
my_partitions)
# Lets add another topic
listener1.reset()
listener2.reset()
my_topic2 = "another-autocreate-pattern-2"
yield from client._wait_on_metadata(my_topic2)
# Wait for group to stabilize
assign1 = yield from listener1.wait_assign()
assign2 = yield from listener2.wait_assign()
# We expect 2 partitons for autocreated topics
my_partitions = set([
TopicPartition(my_topic, 0), TopicPartition(my_topic, 1),
TopicPartition(my_topic2, 0), TopicPartition(my_topic2, 1)])
self.assertEqual(assign1 | assign2, my_partitions)
self.assertEqual(
consumer1.assignment() | consumer2.assignment(),
my_partitions)
yield from consumer1.stop()
yield from consumer2.stop()
yield from client.close()
def test_rebalance_listener_with_coroutines(self):
yield from self.send_messages(0, list(range(0, 10)))
yield from self.send_messages(1, list(range(10, 20)))
main_self = self
class SimpleRebalanceListener(ConsumerRebalanceListener):
def __init__(self, consumer):
self.consumer = consumer
self.revoke_mock = mock.Mock()
self.assign_mock = mock.Mock()
@asyncio.coroutine
def on_partitions_revoked(self, revoked):
self.revoke_mock(revoked)
# If this commit would fail we will end up with wrong msgs
# eturned in test below
yield from self.consumer.commit()
# Confirm that coordinator is actually waiting for callback to
# complete
yield from asyncio.sleep(0.2, loop=main_self.loop)
main_self.assertTrue(
self.consumer._coordinator.needs_join_prepare)
@asyncio.coroutine
def on_partitions_assigned(self, assigned):
self.assign_mock(assigned)
# Confirm that coordinator is actually waiting for callback to
# complete
yield from asyncio.sleep(0.2, loop=main_self.loop)
main_self.assertFalse(
self.consumer._coordinator.needs_join_prepare)
tp0 = TopicPartition(self.topic, 0)
tp1 = TopicPartition(self.topic, 1)
consumer1 = AIOKafkaConsumer(
loop=self.loop,
group_id="test_rebalance_listener_with_coroutines",
bootstrap_servers=self.hosts,
enable_auto_commit=False,
auto_offset_reset="earliest")
listener1 = SimpleRebalanceListener(consumer1)
consumer1.subscribe([self.topic], listener=listener1)
yield from consumer1.start()
self.add_cleanup(consumer1.stop)
msg = yield from consumer1.getone(tp0)
self.assertEqual(msg.value, b"0")
msg = yield from consumer1.getone(tp1)
self.assertEqual(msg.value, b"10")
listener1.revoke_mock.assert_called_with(set([]))
listener1.assign_mock.assert_called_with(set([tp0, tp1]))
# By adding a 2nd consumer we trigger rebalance
consumer2 = AIOKafkaConsumer(
loop=self.loop,
group_id="test_rebalance_listener_with_coroutines",
bootstrap_servers=self.hosts,
enable_auto_commit=False,
auto_offset_reset="earliest")
listener2 = SimpleRebalanceListener(consumer2)
consumer2.subscribe([self.topic], listener=listener2)
yield from consumer2.start()
self.add_cleanup(consumer2.stop)
msg1 = yield from consumer1.getone()
msg2 = yield from consumer2.getone()
# We can't predict the assignment in test
if consumer1.assignment() == set([tp1]):
msg1, msg2 = msg2, msg1
c1_assignment = set([tp1])
c2_assignment = set([tp0])
else:
c1_assignment = set([tp0])
c2_assignment = set([tp1])
self.assertEqual(msg1.value, b"1")
self.assertEqual(msg2.value, b"11")
listener1.revoke_mock.assert_called_with(set([tp0, tp1]))
self.assertEqual(listener1.revoke_mock.call_count, 2)
listener1.assign_mock.assert_called_with(c1_assignment)
self.assertEqual(listener1.assign_mock.call_count, 2)
listener2.revoke_mock.assert_called_with(set([]))
self.assertEqual(listener2.revoke_mock.call_count, 1)
listener2.assign_mock.assert_called_with(c2_assignment)
self.assertEqual(listener2.assign_mock.call_count, 1)
