如何模拟 psycopg2 光标对象？

Question

我在 Python2 中有这个代码段：

 def super_cool_method():
    con = psycopg2.connect(**connection_stuff)
    cur = con.cursor(cursor_factory=DictCursor)
    cur.execute("Super duper SQL query")
    rows = cur.fetchall()

    for row in rows:
        # do some data manipulation on row
    return rows

我想为此编写一些单元测试。我想知道如何使用 mock.patch 来修补游标和连接变量，以便它们返回一组假数据？我已经为我的单元测试尝试了以下代码段，但无济于事：

 @mock.patch("psycopg2.connect")
@mock.patch("psycopg2.extensions.cursor.fetchall")
def test_super_awesome_stuff(self, a, b):
    testing = super_cool_method()

但我似乎收到以下错误：

 TypeError: can't set attributes of built-in/extension type 'psycopg2.extensions.cursor'

原文由 zyshara 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

由于游标是 con.cursor 的返回值，你只需要模拟连接，然后适当地配置它。例如，

 query_result = [("field1a", "field2a"), ("field1b", "field2b")]
with mock.patch('psycopg2.connect') as mock_connect:
    mock_connect.cursor.return_value.fetchall.return_value = query_result
    super_cool_method()

原文由 chepner 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

您有一系列链式调用，每个调用都返回一个新对象。如果您 只 模拟 psycopg2.connect() 调用，您可以通过 .return_value 属性遵循该调用链（每个生成模拟对象），这些属性引用返回的模拟以进行此类调用：

 @mock.patch("psycopg2.connect")
def test_super_awesome_stuff(self, mock_connect):
    expected = [['fake', 'row', 1], ['fake', 'row', 2]]

    mock_con = mock_connect.return_value  # result of psycopg2.connect(**connection_stuff)
    mock_cur = mock_con.cursor.return_value  # result of con.cursor(cursor_factory=DictCursor)
    mock_cur.fetchall.return_value = expected  # return this when calling cur.fetchall()

    result = super_cool_method()
    self.assertEqual(result, expected)

因为您持有模拟 connect 函数的引用，以及模拟连接和游标对象，您还可以断言它们是否被正确调用：

 mock_connect.assert_called_with(**connection_stuff)
mock_con.cursor.asset_called_with(cursor_factory=DictCursor)
mock_cur.execute.assert_called_with("Super duper SQL query")

如果您不需要测试这些，您可以将 return_value 引用直接链接到 cursor() 调用连接对象的结果：

 @mock.patch("psycopg2.connect")
def test_super_awesome_stuff(self, mock_connect):
    expected = [['fake', 'row', 1], ['fake', 'row' 2]]
    mock_connect.return_value.cursor.return_value.fetchall.return_value = expected

    result = super_cool_method()
    self.assertEqual(result, expected)

请注意，如果您使用连接作为 上下文管理器 来自动提交事务 ，并且 您使用 as 将 __enter__() 返回的对象绑定到一个新名称（因此 with psycopg2.connect(...) as conn: # ... ) 那么你需要在调用链中注入一个额外的 __enter__.return_value ：

 mock_con_cm = mock_connect.return_value  # result of psycopg2.connect(**connection_stuff)
mock_con = mock_con_cm.__enter__.return_value  # object assigned to con in with ... as con
mock_cur = mock_con.cursor.return_value  # result of con.cursor(cursor_factory=DictCursor)
mock_cur.fetchall.return_value = expected  # return this when calling cur.fetchall()

这同样适用于 with conn.cursor() as cursor: conn.cursor.return_value.__enter__.return_value 对象被分配给 as 目标。

原文由 Martijn Pieters 发布，翻译遵循 CC BY-SA 4.0 许可协议