asynctest: 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

About this issue

  • Original URL
  • State: closed
  • Created 8 years ago
  • Reactions: 29
  • Comments: 42 (16 by maintainers)

Most upvoted comments

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

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:

This issue has been fixed in the 0.11.0 release which has just been pushed to pypi.

Feedback is more than welcome! If you see anything buggy or that can be improved, open a new bug!

I’ll review my work and list what’s missing by the end of the week. Exception raised in the context manager are not propagated correctly yet: I need to write the tests to ensure it works.

I also would like to add an API which helps manipulating those mocks, but I’ll probably do that in an other PR.