如何用Unittest Mock作为上下文管理器的数据库连接并修复Mock失效问题

问题根源 #

你出错的核心原因是上下文管理器的魔术方法配置错误。with ... as cnxn 语法会自动调用上下文管理器的__enter__方法，取该方法的返回值赋值给as后面的cnxn变量。你原来的代码直接把__enter__属性赋值为你定义的cnxn对象，执行时Python会把__enter__作为方法调用，相当于执行cnxn()，返回的是一个全新的默认MagicMock对象，不是你预先配置过cursor的那个cnxn实例，所以你的cursor配置完全没生效。

修复方案 #

方案1：修正上下文管理器的Mock配置

仅需修改一行代码，指定__enter__方法的返回值为你预先配置好的cnxn对象即可：

@patch("pyodbc.connect")
class testStatCollector(unittest.TestCase):
    def test_initial_values_setted(self, mock_connect):
        cursor = MagicMock(name="my_cursor")
        cursor.fetchall.return_value = [("2", "A", "B")]
        cnxn = MagicMock(name="my_cnxn_mfk")
        cnxn.cursor.return_value = cursor
        # 仅修改这一行即可
        mock_connect.return_value.__enter__.return_value = cnxn
        # 注意检查类名是否和你实际定义的一致，你的示例源码类名是StatCollector，不是PollsStatCollector
        self.test_class = StatCollector(IVR)
        self.assertEqual(
            self.test_class.polls, {"2": Poll("2", "A", "B")}
        )
        self.assertIsInstance(self.test_class.period_start_time, datetime)

方案2：直接Mock私有方法（更简洁）

你的核心需求是跳过数据库逻辑直接给self.polls赋值，完全不需要逐层Mock数据库连接、游标，直接Mock目标私有方法即可。注意双下划线开头的私有方法会被Python做名称修饰，实际属性名格式为_类名__方法名：

@patch.object(StatCollector, "_StatCollector__get_polls")
class testStatCollector(unittest.TestCase):
    def test_initial_values_setted(self, mock_get_polls):
        # 直接设置私有方法返回预设值即可
        mock_get_polls.return_value = {"2": Poll("2", "A", "B")}
        self.test_class = StatCollector(IVR)
        self.assertEqual(
            self.test_class.polls, {"2": Poll("2", "A", "B")}
        )
        self.assertIsInstance(self.test_class.period_start_time, datetime)

该方案代码更简洁，也不会依赖内部数据库调用的实现细节，只要__get_polls的返回格式符合要求就能正常运行。

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