Python unit test mock object within contextmanager

Question

I am trying to mock an object defined inside a ContextManager, and cannot figure it out. Paraphrasing code to test:

def my_func():
  # psycopg2 connection and cursor
  with conn.connection.cursor() as cursor:
     data = cur.execute(<some_sql>)
     writer = csv.writer(<some_open_file>)
     writer.writerows(data)

Now I need to unit test this

@patch_import('csv.writer')
def test_my_func(mock_csv_writer):

    mock_cursor = MagicMock(return_value=MagicMock())
    mock_cursor.return_value.__enter__.return_value = MagicMock()
    mock_conn = MagicMock(connection=MagicMock(cursor=mock_cursor))
    mock_ctx_entry = mock_cursor.return_value.__enter__()

    with mock.patch.object(mock_ctx_entry, "execute") as mock_execute,\
       mock.patch.object(mock_csv_writer, "writerows") as mock_write:
           # Invoke code that calls function
           mock_execute.assert_called
           mock_write.assert_called

However, this does not work because the writer inside the cursor context manager is not the same as the one patched. So how can one mock objects defined inside a context manager?

In fact any documentation on chaining the enter for context managers and their return values would be helpful too.

Python unit test mock object within contextmanager

0 Answers0