如何在Python中模拟嵌套上下文管理器?

2024-04-25 17:13:38 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个Python函数,它使用pgdb连接到数据库并执行查询。你知道吗

    def execute_query(credentials, query, data):
        with pgdb.connect(credentials) as conn:
            with conn.cursor() as cur:
                cur.execute(query, data)

我正试图为这个函数编写一个单元测试,但是在让mock工作时遇到了问题。我遵循了this post的建议,但在尝试模拟嵌套上下文管理器时遇到了问题。你知道吗

最后一个断言失败,因为out是None。你知道吗

@mock.patch("pgdb.connect")
def test_execute_query(self, mock_connect):
    query_result = "test"

    cursor_mock = MagicMock()
    cursor_mock.cursor.__enter__.return_value.execute = query_result

    mock_connect.return_value.__enter__.return_value = cursor_mock

    out = execute_query(None, None, None)
    mock_connect.assert_called_once_with()
    mock_connect.return_value.__enter__.assert_called_once()

    cursor_mock.cursor.assert_called_once()
    cursor_mock.cursor.__enter__.assert_called_once()

    self.assertEquals(out, query_result)

有没有人对我如何让这个工作有什么建议?你知道吗


Tags: noneexecutereturnvalueconnectwithassertout