Python中无法Mock Kafka Producer的produce方法问题求助

我有一段简易代码需要测试：当处理流程无法从事件中提取信息时，应生成一条Kafka消息。我希望通过Mock Kafka Producer来测试该场景，统计produce方法的调用次数。

代码示例如下：

from confluent_kafka import Producer
import os

class MyClass():
    def __init__(self):
        self.producer = Producer({'bootstrap.servers': os.getenv("KAFKA_HOST")})
    
    def my_method(self, value):
        if value <= 0:
            self.producer.produce("TOPIC", {"value": value})

    def process(self, target):
        ... # some magic
        my_method(value) 

我编写的测试代码如下：

import unittest

class TestMyClass(unittest.TestCase):
    def setUpClass(cls) -> None:
        cls.clazz = MyClass()

    @unittest.mock.patch("my.path.to.mymodule.Producer.produce")
    def test_my_class_produces_kafka_msg(self, mock_produce):
        self.clazz.process("dummy")
        assert mock_produce.call_count == 1

但测试报错：TypeError: cannot set 'produce' attribute of immutable type 'cimpl.Producer'。我推测该方法属于不可变类型，想询问是否有绕过方案，或是更优选择为封装Kafka Producer调用的包装方法后再Mock？

临时绕过方案：Mock整个Producer类 #

这个错误的原因是confluent_kafka.Producer是C扩展实现的类，它的方法属于不可变类型，无法直接Mock单个方法。可以改为Mock整个Producer类：

import unittest
from unittest.mock import patch, MagicMock

class TestMyClass(unittest.TestCase):
    @patch("my.path.to.mymodule.Producer")
    def test_my_class_produces_kafka_msg(self, mock_producer_cls):
        # 生成Mock Producer实例
        mock_producer = MagicMock()
        mock_producer_cls.return_value = mock_producer
        
        # 需在Mock生效后初始化MyClass
        clazz = MyClass()
        clazz.process("dummy")
        
        # 验证produce方法调用次数与参数
        mock_producer.produce.assert_called_once_with("TOPIC", {"value": ...})  # 替换为实际预期的value值
        assert mock_producer.produce.call_count == 1

注意：要移除setUpClass中提前初始化MyClass的逻辑，必须在Mock生效后再创建MyClass实例，否则Mock不会生效。

更优方案：封装Kafka Producer调用 #

从代码可维护性和测试友好性角度，封装Producer操作是更好的长期方案。通过将Kafka消息发送逻辑封装为独立模块，不仅能轻松Mock，还能降低业务代码与Kafka客户端的耦合：

步骤1：重构业务代码 #

from confluent_kafka import Producer
import os

class KafkaProducerWrapper:
    def __init__(self):
        self.producer = Producer({'bootstrap.servers': os.getenv("KAFKA_HOST")})
    
    def send_message(self, topic, value):
        self.producer.produce(topic, value)

class MyClass():
    def __init__(self, kafka_producer=None):
        # 依赖注入，测试时可传入Mock实例
        self.kafka_producer = kafka_producer or KafkaProducerWrapper()
    
    def my_method(self, value):
        if value <= 0:
            self.kafka_producer.send_message("TOPIC", {"value": value})

    def process(self, target):
        ... # 业务逻辑
        self.my_method(value) 

步骤2：编写测试代码 #

import unittest
from unittest.mock import MagicMock

class TestMyClass(unittest.TestCase):
    def test_my_class_produces_kafka_msg(self):
        # 创建Mock的Kafka封装类实例
        mock_kafka_wrapper = MagicMock()
        
        # 注入Mock实例初始化MyClass
        clazz = MyClass(kafka_producer=mock_kafka_wrapper)
        clazz.process("dummy")
        
        # 验证消息发送方法的调用
        mock_kafka_wrapper.send_message.assert_called_once_with("TOPIC", {"value": ...})
        assert mock_kafka_wrapper.send_message.call_count == 1

这种方案的优势在于：后续更换Kafka客户端、调整消息发送逻辑（比如添加重试、序列化）时，只需修改封装层，业务代码无需改动；同时测试逻辑更清晰，Mock的是自己定义的接口，不会受第三方库的限制。

内容的提问来源于stack exchange，提问作者Frank