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