Mocking python kafka producer

Question

I have a file with a function which creates a Kafka Producer and publishes some message on to Kafka topic.

def publishToKafkaTopic(value_str):

    producer = KafkaProducer(
        bootstrap_servers=KAFKA_SERVERS,
        value_serializer=lambda x: dumps(x).encode("utf-8"),
    )


    try:
        ack = producer.send(TOPIC_NAME, key=module, value=value_str)
        metadata = ack.get()
        return "Success"
    except KafkaError as kafka_error:
        GLOBAL_LOGGER.error(
            "Failed to publish record on to Kafka broker with error %s", kafka_error
        )


    return "Failed"

Now, I want to test my producer by mocking the KafkaProducer and producer.send(). How do I do that??

score 1 · Answer 1 · answered Feb 03 '21 at 23:58

Note: you need to import your module within the test.

You can test if your producer got called in the first test below, also you can test what args you passed to it in the second test:

from unittest.mock import patch

@patch('kafka.KafkaProducer')
def test_produce_message(KafkaProducerMock):
    publishToKafkaTopic('data')
    KafkaProducerMock.send.assert_called
@patch('kafka.KafkaProducer.send')
def test_vertify_produce_message(KafkaProducerMock):
    publishToKafkaTopic({'key1':'value1', 'key2': 'value2'})
    args = KafkaProducerMock.call_args
    assert(args[0] == (TOPIC_NAME,))
    assert(args[2] == {'value': {'key1':'value1', 'key2': 'value2'}})

score 0 · Answer 2 · answered Nov 28 '20 at 17:16

Replace the KafkaProducer class with a MagicMock() instance in the module where publishToKafkaTopic() is defined and then check that it is called correctly:

import module_importing_kafka_producer

kafka_producer_class_mock = MagicMock()
# replace the imported class
module_importing_kafka_producer.KafkaProducer = kafka_producer_class_mock

# the method to test
module_importing_kafka_producer.publishToKafkaTopic()

# Checks:
# created a KafkaProducer instance
assert kafka_producer_class_mock.call_count == 1

# calls send() on the KafkaProducer instance
kp_instance = run_producer_mock.return_value
assert kp_instance.send.call_count == 1

Beware: untested, but this worked:

def test_whatever():
    mock = MagicMock()

    # simulate the to be tested method and what it does
    mock().send('x')

    # created an instance
    assert mock.call_count == 1
    # called send() on the instance
    assert mock.return_value.send.call_count == 1
    assert mock.return_value.send.call_args.args[0] == 'x'

Mocking python kafka producer

2 Answers2