class test_Producer:
@pytest.fixture()
def producer(self, *, app, _producer):
producer = Producer(app.transport)
producer._producer = _producer
return producer
@pytest.fixture()
def _producer(self):
return Mock(
name='AIOKafkaProducer',
autospec=aiokafka.AIOKafkaProducer,
start=AsyncMock(),
stop=AsyncMock(),
begin_transaction=AsyncMock(),
commit_transaction=AsyncMock(),
abort_transaction=AsyncMock(),
stop_transaction=AsyncMock(),
maybe_begin_transaction=AsyncMock(),
commit=AsyncMock(),
send=AsyncMock(),
flush=AsyncMock(),
)
@pytest.mark.conf(producer_partitioner=my_partitioner)
def test_producer__uses_custom_partitioner(self, *, producer):
assert producer.partitioner is my_partitioner
@pytest.mark.asyncio
async def test_begin_transaction(self, *, producer, _producer):
await producer.begin_transaction('tid')
_producer.begin_transaction.assert_called_once_with('tid')
@pytest.mark.asyncio
async def test_commit_transaction(self, *, producer, _producer):
await producer.commit_transaction('tid')
_producer.commit_transaction.assert_called_once_with('tid')
@pytest.mark.asyncio
async def test_abort_transaction(self, *, producer, _producer):
await producer.abort_transaction('tid')
_producer.abort_transaction.assert_called_once_with('tid')
@pytest.mark.asyncio
async def test_stop_transaction(self, *, producer, _producer):
await producer.stop_transaction('tid')
_producer.stop_transaction.assert_called_once_with('tid')
@pytest.mark.asyncio
async def test_maybe_begin_transaction(self, *, producer, _producer):
await producer.maybe_begin_transaction('tid')
_producer.maybe_begin_transaction.assert_called_once_with('tid')
@pytest.mark.asyncio
async def test_commit_transactions(self, *, producer, _producer):
tid_to_offset_map = {'t1': {TP1: 1001}, 't2': {TP2: 2002}}
await producer.commit_transactions(
tid_to_offset_map, 'group_id', start_new_transaction=False)
_producer.commit.assert_called_once_with(
tid_to_offset_map, 'group_id', start_new_transaction=False)
def test__settings_extra(self, *, producer, app):
app.in_transaction = True
assert producer._settings_extra() == {'acks': 'all'}
app.in_transaction = False
assert producer._settings_extra() == {}
def test__new_producer(self, *, producer):
self.assert_new_producer(producer)
@pytest.mark.parametrize('expected_args', [
pytest.param({'api_version': '0.10'},
marks=pytest.mark.conf(
producer_api_version='0.10')),
pytest.param({'acks': 'all'},
marks=pytest.mark.conf(
producer_acks='all')),
pytest.param({'bootstrap_servers': ['a:9092', 'b:9092']},
marks=pytest.mark.conf(
broker='kafka://a:9092;b:9092')),
pytest.param({'client_id': 'foo'},
marks=pytest.mark.conf(
broker_client_id='foo')),
pytest.param({'compression_type': 'snappy'},
marks=pytest.mark.conf(
producer_compression_type='snappy')),
pytest.param({'linger_ms': 9345},
marks=pytest.mark.conf(
producer_linger_ms=9345)),
pytest.param({'max_batch_size': 41223},
marks=pytest.mark.conf(
producer_max_batch_size=41223)),
pytest.param({'max_request_size': 183831},
marks=pytest.mark.conf(
producer_max_request_size=183831)),
pytest.param({'request_timeout_ms': 1234134000},
marks=pytest.mark.conf(
producer_request_timeout=1234134)),
pytest.param(
{'security_protocol': 'SASL_PLAINTEXT',
'sasl_mechanism': 'PLAIN',
'sasl_plain_username': '******',
'sasl_plain_password': '******',
'ssl_context': None},
marks=pytest.mark.conf(
broker_credentials=auth.SASLCredentials(
username='******',
password='******',
mechanism='PLAIN'),
),
),
])
def test__new_producer__using_settings(self, expected_args, *,
app, producer):
self.assert_new_producer(producer, **expected_args)
def assert_new_producer(self, producer,
acks=-1,
api_version='auto',
bootstrap_servers=['localhost:9092'], # noqa,
client_id=f'faust-{faust.__version__}',
compression_type=None,
linger_ms=0,
max_batch_size=16384,
max_request_size=1000000,
request_timeout_ms=1200000,
security_protocol='PLAINTEXT',
**kwargs):
with patch('aiokafka.AIOKafkaProducer') as AIOKafkaProducer:
p = producer._new_producer()
assert p is AIOKafkaProducer.return_value
AIOKafkaProducer.assert_called_once_with(
acks=acks,
api_version=api_version,
bootstrap_servers=bootstrap_servers,
client_id=client_id,
compression_type=compression_type,
linger_ms=linger_ms,
max_batch_size=max_batch_size,
max_request_size=max_request_size,
request_timeout_ms=request_timeout_ms,
security_protocol=security_protocol,
loop=producer.loop,
partitioner=producer.partitioner,
on_irrecoverable_error=producer._on_irrecoverable_error,
**kwargs,
)
def test__new_producer__default(self, *, producer):
p = producer._new_producer()
assert isinstance(p, aiokafka.AIOKafkaProducer)
def test__new_producer__in_transaction(self, *, producer):
producer.app.in_transaction = True
p = producer._new_producer()
assert isinstance(p, aiokafka.MultiTXNProducer)
def test__producer_type(self, *, producer, app):
app.in_transaction = True
assert producer._producer_type is aiokafka.MultiTXNProducer
app.in_transaction = False
assert producer._producer_type is aiokafka.AIOKafkaProducer
@pytest.mark.asyncio
async def test__on_irrecoverable_error(self, *, producer):
exc = KeyError()
producer.crash = AsyncMock()
app = producer.transport.app
app.consumer = None
await producer._on_irrecoverable_error(exc)
producer.crash.assert_called_once_with(exc)
app.consumer = Mock(name='consumer')
app.consumer.crash = AsyncMock()
await producer._on_irrecoverable_error(exc)
app.consumer.crash.assert_called_once_with(exc)
@pytest.mark.asyncio
async def test_create_topic(self, *, producer, _producer):
producer.transport = Mock(
_create_topic=AsyncMock(),
)
await producer.create_topic(
'foo', 100, 3,
config={'x': 'y'},
timeout=30.3,
retention=300.3,
compacting=True,
deleting=True,
ensure_created=True,
)
producer.transport._create_topic.coro.assert_called_once_with(
producer,
_producer.client,
'foo',
100,
3,
config={'x': 'y'},
timeout=int(30.3 * 1000.0),
retention=int(300.3 * 1000.0),
compacting=True,
deleting=True,
ensure_created=True,
)
def test__ensure_producer(self, *, producer, _producer):
assert producer._ensure_producer() is _producer
producer._producer = None
with pytest.raises(NotReady):
producer._ensure_producer()
@pytest.mark.asyncio
async def test_on_start(self, *, producer, loop):
producer._new_producer = Mock(
name='_new_producer',
return_value=Mock(
start=AsyncMock(),
),
)
_producer = producer._new_producer.return_value
producer.beacon = Mock()
producer._last_batch = loop.time()
await producer.on_start()
assert producer._producer is _producer
producer._new_producer.assert_called_once_with()
producer.beacon.add.assert_called_once_with(_producer)
_producer.start.coro.assert_called_once_with()
assert producer._last_batch is None
@pytest.mark.asyncio
async def test_on_stop(self, *, producer, _producer):
await producer.on_stop()
assert producer._producer is None
_producer.stop.assert_called_once_with()
def test_supports_headers__not_ready(self, *, producer):
producer._producer.client = None
with pytest.raises(NotReady):
producer.supports_headers()
@pytest.mark.asyncio
async def test_send(self, producer, _producer):
await producer.send(
'topic', 'k', 'v', 3, 100, {'foo': 'bar'},
transactional_id='tid',
)
_producer.send.assert_called_once_with(
'topic', 'v',
key='k',
partition=3,
timestamp_ms=100 * 1000.0,
headers=[('foo', 'bar')],
transactional_id='tid',
)
@pytest.mark.asyncio
@pytest.mark.conf(producer_api_version='0.10')
async def test_send__request_no_headers(self, producer, _producer):
await producer.send(
'topic', 'k', 'v', 3, 100, {'foo': 'bar'},
transactional_id='tid',
)
_producer.send.assert_called_once_with(
'topic', 'v',
key='k',
partition=3,
timestamp_ms=100 * 1000.0,
headers=None,
transactional_id='tid',
)
@pytest.mark.asyncio
@pytest.mark.conf(producer_api_version='0.11')
async def test_send__kafka011_supports_headers(self, producer, _producer):
await producer.send(
'topic', 'k', 'v', 3, 100, {'foo': 'bar'},
transactional_id='tid',
)
_producer.send.assert_called_once_with(
'topic', 'v',
key='k',
partition=3,
timestamp_ms=100 * 1000.0,
headers=[('foo', 'bar')],
transactional_id='tid',
)
@pytest.mark.asyncio
@pytest.mark.conf(producer_api_version='auto')
async def test_send__auto_passes_headers(self, producer, _producer):
await producer.send(
'topic', 'k', 'v', 3, 100, [('foo', 'bar')],
transactional_id='tid',
)
_producer.send.assert_called_once_with(
'topic', 'v',
key='k',
partition=3,
timestamp_ms=100 * 1000.0,
headers=[('foo', 'bar')],
transactional_id='tid',
)
@pytest.mark.asyncio
async def test_send__no_headers(self, producer, _producer):
await producer.send(
'topic', 'k', 'v', 3, 100, None,
transactional_id='tid',
)
_producer.send.assert_called_once_with(
'topic', 'v',
key='k',
partition=3,
timestamp_ms=100 * 1000.0,
headers=None,
transactional_id='tid',
)
@pytest.mark.asyncio
async def test_send__no_timestamp(self, producer, _producer):
await producer.send(
'topic', 'k', 'v', 3, None, None,
transactional_id='tid',
)
_producer.send.assert_called_once_with(
'topic', 'v',
key='k',
partition=3,
timestamp_ms=None,
headers=None,
transactional_id='tid',
)
@pytest.mark.asyncio
async def test_send__KafkaError(self, producer, _producer):
_producer.send.coro.side_effect = KafkaError()
with pytest.raises(ProducerSendError):
await producer.send(
'topic', 'k', 'v', 3, None, None,
transactional_id='tid',
)
@pytest.mark.asyncio
async def test_send_and_wait(self, producer):
producer.send = AsyncMock(return_value=done_future(done_future()))
await producer.send_and_wait(
'topic', 'k', 'v', 3, 100, [('a', 'b')],
transactional_id='tid')
producer.send.assert_called_once_with(
'topic',
key='k',
value='v',
partition=3,
timestamp=100,
headers=[('a', 'b')],
transactional_id='tid',
)
@pytest.mark.asyncio
async def test_flush(self, *, producer, _producer):
producer._producer = None
await producer.flush()
producer._producer = _producer
await producer.flush()
_producer.flush.assert_called_once_with()
def test_key_partition(self, *, producer, _producer):
x = producer.key_partition('topic', 'k')
assert x == TP('topic', _producer._partition.return_value)
def test_supports_headers(self, *, producer):
producer._producer.client.api_version = (0, 11)
assert producer.supports_headers()
replication=3,
)
assert 'foo' not in transport._topic_waiters
@pytest.mark.parametrize('credentials,ssl_context,expected', [
(None, {}, {
'security_protocol': 'SSL', 'ssl_context': {},
}),
(None, None, {
'security_protocol': 'PLAINTEXT',
}),
(auth.SSLCredentials({}), None, {
'security_protocol': 'SSL', 'ssl_context': {},
}),
(auth.SASLCredentials(username='******', password='******'), None, {
'security_protocol': 'SASL_PLAINTEXT',
'sasl_mechanism': 'PLAIN',
'sasl_plain_username': '******',
'sasl_plain_password': '******',
'ssl_context': None,
}),
(auth.GSSAPICredentials(kerberos_service_name='service',
kerberos_domain_name='moo'), None, {
'security_protocol': 'SASL_PLAINTEXT',
'sasl_mechanism': 'GSSAPI',
'sasl_kerberos_service_name': 'service',
'sasl_kerberos_domain_name': 'moo',
'ssl_context': None,
}),
])
