def test_compressed_buffer_close(compression):
records = MessageSetBuffer(io.BytesIO(), 100000, compression_type=compression)
orig_msg = Message(b'foobar')
records.append(1234, orig_msg)
records.close()
msgset = MessageSet.decode(records.buffer())
assert len(msgset) == 1
(offset, size, msg) = msgset[0]
assert offset == 0
assert msg.is_compressed()
msgset = msg.decompress()
(offset, size, msg) = msgset[0]
assert not msg.is_compressed()
assert offset == 1234
assert msg == orig_msg
# Closing again should work fine
records.close()
msgset = MessageSet.decode(records.buffer())
assert len(msgset) == 1
(offset, size, msg) = msgset[0]
assert offset == 0
assert msg.is_compressed()
msgset = msg.decompress()
(offset, size, msg) = msgset[0]
assert not msg.is_compressed()
assert offset == 1234
assert msg == orig_msg
def test_produce_request(sender, mocker, api_version, produce_version):
sender.config['api_version'] = api_version
tp = TopicPartition('foo', 0)
records = MessageSetBuffer(io.BytesIO(), 100000)
batch = RecordBatch(tp, records)
produce_request = sender._produce_request(0, 0, 0, [batch])
assert isinstance(produce_request, ProduceRequest[produce_version])
def test_buffer_close():
records = MessageSetBuffer(io.BytesIO(), 100000)
orig_msg = Message(b'foobar')
records.append(1234, orig_msg)
records.close()
msgset = MessageSet.decode(records.buffer())
assert len(msgset) == 1
(offset, size, msg) = msgset[0]
assert offset == 1234
assert msg == orig_msg
# Closing again should work fine
records.close()
msgset = MessageSet.decode(records.buffer())
assert len(msgset) == 1
(offset, size, msg) = msgset[0]
assert offset == 1234
assert msg == orig_msg
def add_message(self, tp, key, value, timeout):
"""Add message to batch by topic-partition
If batch is already full this method waits (`ttl` seconds maximum)
until batch is drained by send task
"""
if self._closed:
# this can happen when producer is closing but try to send some
# messages in async task
raise ProducerClosed()
batch = self._batches.get(tp)
if not batch:
message_set_buffer = MessageSetBuffer(io.BytesIO(),
self._batch_size,
self._compression_type)
batch = MessageBatch(tp, message_set_buffer, self._batch_ttl,
self._loop)
self._batches[tp] = batch
if not self._wait_data_future.done():
# Wakeup sender task if it waits for data
self._wait_data_future.set_result(None)
future = batch.append(key, value)
if future is None:
# Batch is full, can't append data atm,
# waiting until batch per topic-partition is drained
start = self._loop.time()
yield from asyncio.wait([batch.wait_drain()],
timeout=timeout,
loop=self._loop)
timeout -= self._loop.time() - start
if timeout <= 0:
raise KafkaTimeoutError()
return (yield from self.add_message(tp, key, value, timeout))
return future
