web-dev-qa-db-ja.com

psycopg2カーソルオブジェクトをモックする方法は?

私は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'
19
zyshara

カーソルは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()
10
chepner

一連の連鎖呼び出しがあり、それぞれが新しいオブジェクトを返します。 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として使用してトランザクションを自動的にコミットする場合および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ターゲットに割り当てられます。

33
Martijn Pieters