Имитация пользовательского контекстного менеджера внутри функции

У меня есть собственный менеджер контекста, определенный как:

# ~/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()

В функции с именем get_data я передаю результат диспетчера контекста conn в качестве аргумента другой функции 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

Я хочу проверить, чтобы убедиться, что get_data_for_id вызывается с определенными аргументами, когда я вызываю 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])

Теперь по какой-то причине conn, возвращаемый из query.get_db_connection в main.py, является объектом соединения, когда я ожидаю, что это будет объект Mock(), потому что я исправил его. Я получаю следующие ошибки:

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)]

Поскольку mock_conn уже является объектом Mock(), другие решения по переполнению стека, которые указывают на решения путем изменения .return_value диспетчера контекста, не применяются.

Как я могу сделать оба объекта conn одинаковыми Mock()?


person scrollout    schedule 26.01.2021    source источник
comment
Вы импортируете query в main с чем-то вроде from db import query, поэтому вам нужно исправить этот экземпляр query, например. @patch('main.query.get_db_connection') — см. где исправить.   -  person MrBean Bremen    schedule 26.01.2021
comment
Менеджер контекста теперь вызывается с аргументами: [call(<MagicMock name='get_db_connection().__enter__()' id='261880080'>, 1), call(<MagicMock name='get_db_connection().__enter__()' id='261880080'>, 2),...]   -  person scrollout    schedule 27.01.2021
comment
Мне пришлось использовать mock_conn.return_value.__enter__.return_value в дополнение к вашему решению   -  person scrollout    schedule 27.01.2021
comment
Правильно, вы всегда должны делать это для доступа к диспетчеру контекста, забыл упомянуть об этом .   -  person MrBean Bremen    schedule 27.01.2021