python - Mocking a custom context manager within a function

Question

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:UsersscrolloutAppDataLocalProgramsPythonPython38libunittestmock.py", line 1325, in patched
    return func(*newargs, **newkeywargs)
  File "pkgest.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()?

question from:https://stackoverflow.com/questions/65907143/mocking-a-custom-context-manager-within-a-function

Answer 1

Waitting for answers