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

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

[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

[thehesiod]

nice work @Martiusweb on a real fix and @Galbar on work-around, this thread inspired us to do a simpler fix:

class MagicMockContext(MagicMock):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        type(self).__aenter__ = CoroutineMock(return_value=MagicMock())
        type(self).__aexit__ = CoroutineMock(return_value=MagicMock())

with usage:

with asynctest.mock.patch('dummy.Dummy', new_callable=MagicMockContext) as mockClientSession: