I have a Python function that connects to a database using pgdb and executes a query.
def execute_query(credentials, query, data):
with pgdb.connect(credentials) as conn:
with conn.cursor() as cur:
cur.execute(query, data)
I am trying to write a unit test for this function but am having trouble getting mocking to work. I followed the suggestions from this post but am having issues when trying to mock the nested context manager.
The last assert is failing because out
is None.
@mock.patch("pgdb.connect")
def test_execute_query(self, mock_connect):
query_result = "test"
cursor_mock = MagicMock()
cursor_mock.cursor.__enter__.return_value.execute = query_result
mock_connect.return_value.__enter__.return_value = cursor_mock
out = execute_query(None, None, None)
mock_connect.assert_called_once_with()
mock_connect.return_value.__enter__.assert_called_once()
cursor_mock.cursor.assert_called_once()
cursor_mock.cursor.__enter__.assert_called_once()
self.assertEquals(out, query_result)
Does anyone have any suggestions on how I can get this to work?