def test_consume_one_thousand_of_messages(self):
n_messages = 1000
consumed_messages = asyncio.Queue(maxsize=n_messages, loop=self.loop)
async def consume_messages():
async for message in self.stream_consumer:
consumed_messages.put_nowait(message)
self.stream_consumer.start()
produce_to_kafka(self.test_topic, self.test_message, number=1000)
asyncio.ensure_future(consume_messages(), loop=self.loop)
async def wait_for_messages():
while True:
await asyncio.sleep(0.1)
if consumed_messages.qsize() == n_messages:
break
coro = asyncio.wait_for(wait_for_messages(),
timeout=30,
loop=self.loop)
self.loop.run_until_complete(coro)
for _ in range(n_messages):
self.assertEqual(consumed_messages.get_nowait().payload,
self.test_message)
def test_assign_topic_offset(self):
n_messages = 100
timeout = 5.0
produce_to_kafka(self.test_topic, self.test_message, number=n_messages)
seek_offset = 50
self.stream_consumer.start()
self.loop.run_until_complete(
asyncio.wait_for(self.stream_consumer.__anext__(),
timeout=timeout))
self.stream_consumer.assign_topic_offset(topic=self.test_topic,
partition=0,
offset=seek_offset)
async def consume_messages(buffer, limit):
async for message in self.stream_consumer:
buffer.put_nowait(message)
if buffer.qsize() >= limit:
break
buffer = asyncio.Queue(maxsize=n_messages, loop=self.loop)
self.loop.run_until_complete(
asyncio.wait_for(consume_messages(buffer, limit=1),
timeout=timeout))
first_element = buffer.get_nowait()
self.assertEqual(first_element.offset, seek_offset)
def test_consume_one_message_with_key(self):
confirm_message = asyncio.Future(loop=self.loop)
async def consume_messages():
async for message in self.stream_consumer:
confirm_message.set_result(message)
self.stream_consumer.start()
produce_to_kafka(self.test_topic, self.test_message, key=self.test_key)
asyncio.ensure_future(consume_messages(), loop=self.loop)
coro = asyncio.wait_for(confirm_message, timeout=10, loop=self.loop)
self.loop.run_until_complete(coro)
consumed_message = confirm_message.result()
self.assertEqual(consumed_message.payload, self.test_message)
self.assertEqual(consumed_message.topic, self.test_topic)
self.assertEqual(consumed_message.key, self.test_key)
self.assertTrue(isinstance(consumed_message.offset, int))