如何模拟 Python 单元测试的 Kafka Producer 和 producer.send 方法

Question

I am new to Python and trying to write a unit test which involves Kafka. I have a class with a function which calls another function to initialize Kafka Producer and then call producer.send(). I want to mock Kafka Producer for my unit test.

Below is my code and I want to write unit test for produce_kafka_message method.

class KafkaProducerIntface
    
    def __init__(self, topic, .....)
        self.bootstrap_server = bootstrap_server
        self.producer = None
        self.post_topic = topic
    ....

    def produce_kafka_message(self, key, value, headers)

        self.__initialize_producer__(retries=3)
        future = self.producer.send(self.topic, key=key, value=value, headers=headers)
        self.producer.flush()


    def __initialize_producer__(self, retries=3)
        self.producer = KafkaProducer(bootstrap_servers=self.bootstrap_server, acks='all', retries=retries)

Answer 1

You can use mock lib for it and mock the produce_kafka_message method:

with mock.patch.object(KafkaProducerIntface, "produce_kafka_message") as mock_produce_kafka_message:
    ... # do some code here
    mock_produce_kafka_message.assert_called_once_with(key, value, headerss)

You can also do the same with KafkaProducer class it self:

with mock.patch("kafka.path.KafkaProducer") as mock_kafka_producer:
     interface = KafkaProducerIntface(...)
     interface.produce_kafka_message(key, value, headers)
     
     mock_kafka_producer.send.assert_called_once()
     mock_kafka_producer.flush()

BTW. also please check if you want to run flush after the send or just.poll(0):)