如何模拟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'

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

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.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目标。