关闭 contextlib 模块时如何模拟 psycopg2.connect？

我有以下带有 get_data 方法的数据库类。它使用“contextlib”模块的“关闭”来关闭连接。

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'>

注意 ：我无法修改数据库类的实现。

你的问题出在 mock 对象的嵌套方式上。在 get_data 方法中，你使用了两个上下文管理器：一个是 closing(psycopg2.connect(**self.db_details)) ，另一个是 conn.cursor() 。这意味着你需要对你的 mock 对象进行相应的嵌套，以便正确地模拟这些调用。

下面是修正后的测试代码：

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 对象
        mock_conn = mock_connect.return_value
        mock_cursor = mock_conn.cursor.return_value
        mock_cursor.__enter__.return_value.fetchall.return_value = expected_result

        actual_result = db_inst.get_data(query, parameters)

        self.assertEqual(expected_result, actual_result)

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

在这个修正后的代码中：

1. 我们首先创建了 mock_conn 对象，它是 mock_connect.return_value ，表示模拟的连接对象。
2. 然后我们创建了 mock_cursor 对象，它是 mock_conn.cursor.return_value ，表示模拟的游标对象。
3. 最后，我们将 expected_result 设置为 mock_cursor.__enter__.return_value.fetchall.return_value 。这是因为 conn.cursor() 返回一个上下文管理器，我们需要调用 __enter__ 方法获取游标对象，然后才能调用 fetchall 方法。

通过这种方式正确地嵌套 mock 对象，你就可以成功地模拟 psycopg2.connect 和 conn.cursor() 的调用，并验证你的 get_data 方法的行为。