Làm thế nào để giả lập đối tượng con trỏ psycopg2?

Question

Tôi có đoạn mã này trong 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 

mà tôi muốn viết một số unittests cho. Tôi đang tự hỏi làm thế nào để sử dụng mock.patch để vá ra các con trỏ và các biến kết nối để họ trở về một tập hợp giả của dữ liệu? Tôi đã thử các phân đoạn sau của mã cho unittests của tôi, nhưng không có kết quả:

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

Nhưng tôi dường như nhận được lỗi sau:

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

Answer 1

Kể từ khi con trỏ là giá trị trở lại của con.cursor, bạn chỉ cần giả lập kết nối, sau đó cấu hình nó đúng cách. Ví dụ:

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

Answer 2

Bạn có một chuỗi cuộc gọi bị xích, mỗi lần trả về một đối tượng mới. Nếu bạn thử chỉ các psycopg2.connect() cuộc gọi, bạn có thể làm theo mà chuỗi các cuộc gọi (mỗi sản xuất các đối tượng giả) thông qua .return_value thuộc tính, mà tham khảo các mô hình trở lại cho các cuộc gọi như vậy:

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

Bởi vì bạn giữ được tài liệu tham khảo cho mock connect chức năng, cũng như kết nối giả và con trỏ đối tượng bạn có thể sau đó cũng khẳng định nếu họ được gọi một cách chính xác:

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

Nếu bạn không cần phải kiểm tra này, bạn chỉ có thể chuỗi lên tài liệu tham khảo để return_value đi thẳng với kết quả của cuộc gọi cursor() trên đối tượng kết nối:

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