I have a custom context manager defined as:
# ~/db/query.py
@contextmanager
def get_db_connection(conn_string, **kwargs):
try:
conn = pyodbc.connect(conn_string, **kwargs)
yield conn
except Exception as connection_error:
raise ValueError('Could not connect to db.', connection_error) from None
finally:
conn.close()
Within a function called get_data, I pass the result of the context manager, conn, as an argument to another function, get_data_for_id:
# ~/pkg/main.py
def get_data(ids):
dfs = []
with query.get_db_connection(conn_string) as conn:
for id in ids:
df = get_data_for_id(conn, id)
dfs.append(df)
return dfs
I want to test to make sure get_data_for_id is being called with specific arguments when I call get_data:
# ~/pkg/test.py
@patch('db.query.get_db_connection')
@patch('main.get_data_for_id')
def test_get_data(self, mock_query, mock_conn):
ids = [1, 2, 3]
main.get_data(ids)
self.assertEqual(mock_query.call_args_list, [call(mock_conn, x) for x in ids])
Now for some reason, conn returned from query.get_db_connection in main.py is a connection object when I expect it to be a Mock() object because I patched it. I am getting the following errors:
FAIL: test_get_data (__main__.Tests)
----------------------------------------------------------------------
Traceback (most recent call last):
File "C:\Users\scrollout\AppData\Local\Programs\Python\Python38\lib\unittest\mock.py", line 1325, in patched
return func(*newargs, **newkeywargs)
File "\pkg\test.py", line 108, in test_get_data
self.assertEqual(mock_query.call_args_list, [call(mock_conn, x) for x in ids])
AssertionError: [call(<pyodbc.Connection object at 0x000000000FC[126 chars], 3)] != [call(<MagicMock name='get_db_connection' id='26[133 chars], 3)]
Since mock_conn is already a Mock() object, other solutions on stack overflow that point out solutions by changing the context manager's .return_value don't apply.
How can I make both conn objects the same Mock()?
