From adf63e65a0b861cb02a6e8cc1e5b2c2c09e57726 Mon Sep 17 00:00:00 2001 From: Akarys42 Date: Sun, 9 Feb 2020 17:05:09 +0100 Subject: Create an AsyncContextManagerMock mock for testing asynchronous context managers It can be used to test aiohttp request functions, since they are async context managers --- tests/helpers.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) (limited to 'tests/helpers.py') diff --git a/tests/helpers.py b/tests/helpers.py index 5df796c23..6aee8623f 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -127,6 +127,18 @@ class AsyncMock(CustomMockMixin, unittest.mock.MagicMock): return super().__call__(*args, **kwargs) +class AsyncContextManagerMock(unittest.mock.MagicMock): + def __init__(self, return_value: Any): + super().__init__() + self._return_value = return_value + + async def __aenter__(self): + return self._return_value + + async def __aexit__(self, *args): + pass + + class AsyncIteratorMock: """ A class to mock asynchronous iterators. -- cgit v1.2.3