web-dev-qa-db-fra.com

Comment se moquer de l'objet curseur psycopg2?

J'ai ce segment de code en 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

pour lequel j'aimerais écrire quelques tests. Je me demande comment utiliser mock.patch afin de patcher le curseur et les variables de connexion afin qu'ils retournent un faux ensemble de données? J'ai essayé le segment de code suivant pour mes tests, mais en vain:

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

Mais je semble obtenir l'erreur suivante:

TypeError: can't set attributes of built-in/extension type 'psycopg2.extensions.cursor'
pythonpython-2.7unit-testingmockingpsycopg2
19
zyshara

Puisque le curseur est la valeur de retour de con.cursor, il vous suffit de vous moquer de la connexion, puis de la configurer correctement. Par exemple,

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

Vous avez une série d'appels chaînés, chacun renvoyant un nouvel objet. Si vous vous moquez juste de l'appel psycopg2.connect(), vous pouvez suivre cette chaîne d'appels (chacun produisant des objets fictifs) via .return_value attributs, qui font référence à la maquette retournée pour de tels appels:

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

Étant donné que vous conservez les références de la fonction mock connect, ainsi que les objets de connexion et de curseur fictifs, vous pouvez également vérifier s'ils ont été appelés correctement:

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

Si vous n'avez pas besoin de les tester, vous pouvez simplement chaîner les références return_value Pour aller directement au résultat de l'appel de cursor() sur l'objet de connexion:

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

Notez que si vous utilisez la connexion en tant que gestionnaire de contexte pour valider automatiquement la transaction et vous utilisez as pour lier l'objet renvoyé par __enter__() à un nouveau nom (donc with psycopg2.connect(...) as conn: # ...) alors vous besoin d'injecter un __enter__.return_value supplémentaire dans la chaîne d'appel:

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

Il en va de même pour le résultat de with conn.cursor() as cursor:, l'objet conn.cursor.return_value.__enter__.return_value Est affecté à la cible as.

33
Martijn Pieters