[Python] How can I mock multiple context managers and function calls with pytest in async...

I'm writing tests for a python fastApi applciation, however, the issue is not directly related to fastAPI itself. I have a block of code that doesn't seem to be mocked properly, and the offending code is:

async with await self._db.get_pool().getconn() as conn:
async with await conn.cursor() as cursor:
await cursor.execute(...)
result = await cursor.fetchone()
...
return Object(id=result[0]...)


the self._db key is of type DB that is defined elsewhere. the execute call is just a query against the database, and the object type is for example. the return type from the function is being populated by the returned list of values from the execute.

I have the mock setup like this:

mock_db = AsyncMock(spec=DB)
mock_pool = AsyncMock(spec=AsyncConnectionPool)
mock_conn = AsyncMock(spec=AsyncConnection)
mock_cursor = AsyncMock(spec=AsyncCursor)

mock_db.__aenter__.return_value = mock_db
mock_db.get_pool.return_value = mock_pool
mock_pool.getconn.return_value = mock_conn
mock_conn.__aenter__.return_value = mock_conn
mock_conn.cursor.return_value = mock_cursor
mock_cursor.__aenter__.return_value = mock_cursor
mock_cursor.fetchone.return_value = (uuid(), ....)
result = await test_func(...)


when the test runs, we don't get to assertions as the mocks are all messed up.

I know I went wrong somewhere, but where is what I can't tell at the moment.

I've been trying to find documentation specifically for testing this case, but haven't been able to find a clear explanation of how to mock a nested context managers, and an explanation of where I need to mock through the __aenter__ function.

I have also tried to have each mocked value that has __aenter__ to call the return_value property before __aenter__

Continue reading...