test pg connection mock

test_pg_conn_mock.py

#!/usr/bin/env python

from __future__ import print_function

try:
    import mock
except ImportError:
    from unittest import mock
import unittest

import psycopg2


def ex(user_id):
    with psycopg2.connect("") as conn:
        with conn.cursor() as cursor:
            cursor.execute("SELECT first_name FROM users WHERE id = %s", (user_id,))
            return cursor.fetchall()


class TestPatchConn(unittest.TestCase):
    def setUp(self):
        self.user_id = 16

    @mock.patch('psycopg2.connect')
    def test_with_context(self, mock_connect):
        mock_connect().__enter__().cursor().__enter__().fetchall.return_value = ['Stacy']
        ex(self.user_id)
        mock_connect().__enter__().cursor().__enter__().execute.assert_called_with('SELECT first_name FROM users WHERE id = %s', (self.user_id, ))


if __name__ == '__main__':
    unittest.main()