如何Mock结合contextlib.closing的psycopg2.connect？测试断言失败求助

我有一个包含get_data方法的Database类，该方法使用contextlib模块的closing管理数据库连接，代码如下：

class Database:

    def __init__(self) -> None:
        self.db_details = {<connection details>}

    def get_data(self, query, parameters):

        with closing(psycopg2.connect(**self.db_details)) as conn:
            with conn.cursor() as cursor:
                cursor.execute(query, parameters)
                data = cursor.fetchall()

        return data

编写了如下单元测试：

class TestDatabase(unittest.TestCase):

    @patch("psycopg2.connect")
    @patch("Database.__init__")
    def test_get_data_pos(self, mock_init, mock_connect):

        expected_result = [(1, "John", "New York City")]
        query = "SELECT * FROM user_details;"
        parameters = {}

        mock_init.return_value = None

        db_inst = Database()
        db_inst.db_details = {
            <some_dummy_values>
        }

        mock_con_cm = mock_connect.return_value
        mock_conn = mock_con_cm.__enter__.return_value
        mock_cur = mock_conn.cursor.return_value.__enter__.return_value
        mock_cur.fetchall.return_value = expected_result

        actual_result = db_inst.get_data(query, parameters)

        self.assertEqual(expected_result, actual_result)


if __name__ == "__main__":
    unittest.main()

运行后出现断言错误：

AssertionError: [(1, 'John', 'New York City')] != <MagicMock name='connect().cursor().__ent[35 chars]384'>

注意：无法修改Database类的实现，需解决该问题。

解决方案 #

问题出在对closing上下文管理器的模拟逻辑错误。contextlib.closing的__enter__方法会直接返回被包装的对象（即psycopg2.connect返回的连接实例），不需要额外模拟__enter__的返回值。

修改后的测试代码如下：

class TestDatabase(unittest.TestCase):

    @patch("psycopg2.connect")
    @patch("Database.__init__")
    def test_get_data_pos(self, mock_init, mock_connect):

        expected_result = [(1, "John", "New York City")]
        query = "SELECT * FROM user_details;"
        parameters = {}

        mock_init.return_value = None

        db_inst = Database()
        db_inst.db_details = {
            <some_dummy_values>
        }

        # 直接将mock_conn设置为connect的返回值，closing会直接返回它
        mock_conn = mock_connect.return_value
        # 模拟cursor上下文管理器
        mock_cur_cm = mock_conn.cursor.return_value
        mock_cur = mock_cur_cm.__enter__.return_value
        mock_cur.fetchall.return_value = expected_result

        actual_result = db_inst.get_data(query, parameters)

        self.assertEqual(expected_result, actual_result)

原因说明 #

原测试错误地为connect返回值模拟了__enter__方法，但closing上下文管理器仅在上下文退出时调用对象的close方法，进入上下文时直接返回传入的连接对象。之前的模拟路径导致代码实际调用connect().__enter__().cursor()，而真实代码中是connect().cursor()，最终fetchall没有绑定到实际调用的游标对象上，返回了未设置返回值的MagicMock实例。

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