MagicMock doesn't implement `__aenter__` or `__aexit__`

[txomon]

When testing asyncio code, the probability of needing to test async context manager is high.

The current code using MagicMock doesn't implement __aenter__ or __aexit__ methods

[Martiusweb]

Hi,

Thanks for the report, it's on the todo list!

[Gr1N]

+1 it would be really useful to add implementation of __aenter__ and __aexit__ methods.

[txomon]

I add this link as relevant to the conversation http://bugs.python.org/issue26467

[serg666]

Hi, guys! Any news about this issue?

[serg666]

I would like to mock aiopg.sa.create_engine with CoroutineMock, but if I use async-with-style it does not work

import asyncio
import aiopg.sa
import asynctest


mock = asynctest.CoroutineMock()

async def create_engine():
    kwargs = {'echo': True}
    return await aiopg.sa.create_engine(
        'dbname=db user=user password=pass host=127.0.0.1',
        **kwargs
   )

async def go(get_engine):
    engine = await get_engine()
    async with engine.acquire() as conn:
        async with conn.begin():
            r = await conn.execute('SELECT version();')
    return await r.fetchone()

row = asyncio.get_event_loop().run_until_complete(go(create_engine))
print(row)

row = asyncio.get_event_loop().run_until_complete(go(mock))
print(row)

python atest.py 
('PostgreSQL 9.6.1 on x86_64-suse-linux-gnu, compiled by gcc (SUSE Linux) 6.2.1 20161209 [gcc-6-branch revision 243481], 64-bit',)
Traceback (most recent call last):
  File "atest.py", line 25, in <module>
    row = asyncio.get_event_loop().run_until_complete(go(mock))
  File "/usr/lib64/python3.5/asyncio/base_events.py", line 337, in run_until_complete
    return future.result()
  File "/usr/lib64/python3.5/asyncio/futures.py", line 274, in result
    raise self._exception
  File "/usr/lib64/python3.5/asyncio/tasks.py", line 239, in _step
    result = coro.send(None)
  File "atest.py", line 17, in go
    async with engine.acquire() as conn:
AttributeError: __aexit__

but if I do not use async-with-style all works fine

import asyncio
import aiopg.sa
import asynctest


mock = asynctest.CoroutineMock()

async def create_engine():
    kwargs = {'echo': True}
    return await aiopg.sa.create_engine(
        'dbname=wallet user=wallet password=qazwsx host=127.0.0.1',
        **kwargs
   )

async def go(get_engine):
    engine = await get_engine()
    conn = await engine.acquire()
    trans = await conn.begin()
    r = await conn.execute('SELECT version();')
    await trans.commit()
    engine.release(conn)
    return await r.fetchone()

row = asyncio.get_event_loop().run_until_complete(go(create_engine))
print(row)

row = asyncio.get_event_loop().run_until_complete(go(mock))
print(row)

('PostgreSQL 9.6.1 on x86_64-suse-linux-gnu, compiled by gcc (SUSE Linux) 6.2.1 20161209 [gcc-6-branch revision 243481], 64-bit',)
<CoroutineMock name='mock.......' id='139844439568168'>

[Galbar]

We had the same problem mocking aiopg. We ended up with this workaround:

class AsyncContextManagerMock(MagicMock):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        for key in ('aenter_return', 'aexit_return'):
            setattr(self, key,  kwargs[key] if key in kwargs else MagicMock())

    async def __aenter__(self):
        return self.aenter_return

    async def __aexit__(self, *args):
        return self.aexit_return

[Martiusweb]

Hello,

Some news regarding this issue: I'm working on a solution which seems to be working for the basic needs (in the async_with branch). Some cases are yet to be tested, but if some of you want to try it, feedback is welcome.

[serg666]

Hello, @Martiusweb !

Do you mean async_magic branch ?

[Martiusweb]

Yes, I mean async_magic (https://github.com/Martiusweb/asynctest/tree/async_magic).

[serg666]

I have just try async_magic branch with code as below

import asyncio
import aiopg.sa
import asynctest


mock = asynctest.CoroutineMock()

async def create_engine():
    kwargs = {'echo': True}
    return await aiopg.sa.create_engine(
        'dbname=db user=dbuser password=qazwsx host=127.0.0.1',
        **kwargs
   )

async def go(get_engine):
    engine = await get_engine()
    async with engine.acquire() as conn:
        async with conn.begin():
            r = await conn.execute('SELECT version();')
    return await r.fetchone()

row = asyncio.get_event_loop().run_until_complete(go(create_engine))
print(row)

row = asyncio.get_event_loop().run_until_complete(go(mock))
print(row)

I have got the following:

python atest.py
('PostgreSQL 9.6.2 on x86_64-suse-linux-gnu, compiled by gcc (SUSE Linux) 6.3.1 20170202 [gcc-6-branch revision 245119], 64-bit',)
Traceback (most recent call last):
File "atest.py", line 25, in
row = asyncio.get_event_loop().run_until_complete(go(mock))
File "/usr/lib64/python3.6/asyncio/base_events.py", line 466, in run_until_complete
return future.result()
File "atest.py", line 17, in go
async with engine.acquire() as conn:
TypeError: object MagicMock can't be used in 'await' expression

[Martiusweb]

This is because asynctest has no way to know that engine.acquire() is a coroutine, hence it's mocked as a MagicMock rather than a CoroutineMock.

However, it worked before because CoroutineMock returned an other CoroutineMock as the result of its invocation, which is wrong.

You can fix this by doing something like:

engine_mock = asyncio.MagicMock(Engine())
mock = asyncio.CoroutineMock(return_value=engine_mock)