如何模拟 psycopg2 光标对象？答案

【问题标题】：How to mock psycopg2 cursor object?如何模拟 psycopg2 光标对象？
【发布时间】：2016-05-10 15:26:28
【问题描述】：

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

【问题讨论】：

标签： python python-2.7 unit-testing mocking psycopg2

【解决方案1】：

您有一系列链式调用，每个都返回一个新对象。如果您模拟只是 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.called_with(cursor_factory=DictCursor)
mock_cur.execute.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)

请注意，如果您将连接用作context manager to automatically commit the transaction 并且您使用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目标。

【讨论】：

这太棒了，感谢上下文管理器__enter__ 的建议。从哲学的角度来看，这是一种建议的测试方式吗？似乎只是将模拟返回的expected 与expected 进行比较，因此即使super_cool_method() 中的逻辑发生更改，此测试也始终有效——只要语法有效，测试就永远不会中断。这是一个简单的例子来说明这一点，还是在做我错过的事情？我可以看到您可以强制执行 DictCursor 这很有价值，但看不到为什么要测试 fetchall.return_value 并尝试学习。
@Davos 是的，这是一个很好的理念，因为您的单元测试不应该测试 psycopg2 是否正常工作，它应该测试您的代码是否正确使用结果。实际测试中的代码将对fetchall() 的结果做一些事情，这很重要。或者您想测试传递给 cursor.execute() 等的内容。

【解决方案2】：

由于游标是con.cursor的返回值，所以只需要mock连接，然后正确配置即可。例如，

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()

【讨论】：

这对我不起作用。我从 fetchall 得到的结果是 ''
cursor.execute 通常返回 None。为什么要为execute而不是fetchall配置返回值？
因为显然我应该嘲笑fetchall :) 很抱歉造成混乱。
不幸的是仍然不适合我。我已将我的问题作为自己的问题发布：stackoverflow.com/questions/37164891/…
@chepner：实际上，因为connect 和execute 都被调用，所以您需要在此处注入更多return_value 引用。并不是说使用了execute 的返回值，fetchall 从未在该模拟中引用。

【解决方案3】：

以下答案是上述答案的变体。我使用的是django.db.connections 光标对象。

所以下面的代码对我有用

@patch('django.db.connections')
def test_supercool_method(self, mock_connections):
    query_result = [("field1a", "field2a"), ("field1b", "field2b")]
    mock_connections.__getitem__.return_value.cursor.return_value.__enter__.return_value.fetchall.return_value = query_result

    result = supercool_method()
    self.assertIsInstance(result, list)

【讨论】：