Skip to content

Instantly share code, notes, and snippets.

@graphaelli
Created December 9, 2016 03:55
Show Gist options
  • Save graphaelli/906b624c18f77f50da5cd0cd4211c3c8 to your computer and use it in GitHub Desktop.
Save graphaelli/906b624c18f77f50da5cd0cd4211c3c8 to your computer and use it in GitHub Desktop.
test pg connection mock
#!/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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment