bpo-38093: Correctly returns AsyncMock for async subclasses.

Lib/unittest/mock.py

@@ -989,9 +989,9 @@ def _get_child_mock(self, /, **kw):
         _type = type(self)
         if issubclass(_type, MagicMock) and _new_name in _async_method_magics:
             klass = AsyncMock
-        if issubclass(_type, AsyncMockMixin):
-            klass = MagicMock
-        if not issubclass(_type, CallableMixin):
+        elif issubclass(_type, AsyncMockMixin):
+            klass = AsyncMock
+        elif not issubclass(_type, CallableMixin):
             if issubclass(_type, NonCallableMagicMock):
                 klass = MagicMock
             elif issubclass(_type, NonCallableMock) :
@@ -1904,6 +1904,8 @@ def method(self, /, *args, **kw):
     '__str__': lambda self: object.__str__(self),
     '__sizeof__': lambda self: object.__sizeof__(self),
     '__fspath__': lambda self: f"{type(self).__name__}/{self._extract_mock_name()}/{id(self)}",
+    '__aenter__': lambda self: AsyncMockMixin._mock_call(self),
+    '__aexit__': lambda self: AsyncMockMixin._mock_call(self)
 }
 
 _return_values = {

Lib/unittest/test/testmock/testasync.py

@@ -363,22 +363,58 @@ def test_add_side_effect_iterable(self):
 class AsyncContextManagerTest(unittest.TestCase):
 
     class WithAsyncContextManager:
-
         async def __aenter__(self, *args, **kwargs):
             return self
 
         async def __aexit__(self, *args, **kwargs):
             pass
 
-    def test_magic_methods_are_async_mocks(self):
-        mock = MagicMock(self.WithAsyncContextManager())
-        self.assertIsInstance(mock.__aenter__, AsyncMock)
-        self.assertIsInstance(mock.__aexit__, AsyncMock)
+    class WithSyncContextManager:
+        def __enter__(self, *args, **kwargs):
+            return self
+
+        def __exit__(self, *args, **kwargs):
+            pass
+
+    class ProductionCode:
+        # Example real-world(ish) code
+        async def main(self):
+            async with self.session.post('https://python.org') as response:
+                val = await response.json()
+                return val
+
+    def test_async_magic_methods_are_async_mocks_with_magicmock(self):
+        cm_mock = MagicMock(self.WithAsyncContextManager)
+        self.assertIsInstance(cm_mock.__aenter__, AsyncMock)
+        self.assertIsInstance(cm_mock.__aexit__, AsyncMock)
+
+    def test_set_return_value_of_aenter_magicmock(self):
+        pc = self.ProductionCode()
+        pc.session = MagicMock(name='sessionmock')
+        cm = MagicMock(name='magic_cm')
+        response = AsyncMock(name='response')
+        response.json = AsyncMock(return_value={'json':123})
+        cm.__aenter__.return_value = response
+        pc.session.post.return_value = cm
+        result = asyncio.run(pc.main())
+        self.assertEqual(result, {'json':123})
+
+    def test_set_return_value_of_aenter_asyncmock(self):
+        pc = self.ProductionCode()
+        pc.session = MagicMock(name='sessionmock')
+        cm = AsyncMock(name='async_cm')
+        response = AsyncMock(name='response')
+        response.json = AsyncMock(return_value={'json':123})
+        cm.__aenter__.return_value = response
+        pc.session.post.return_value = cm
+        result = asyncio.run(pc.main())
+        self.assertEqual(result, {'json':123})
 
     def test_mock_supports_async_context_manager(self):
         called = False
         instance = self.WithAsyncContextManager()
-        mock_instance = MagicMock(instance)
+        # TODO: Horrible error message if you use a MagicMock here
+        mock_instance = AsyncMock(instance)
 
         async def use_context_manager():
             nonlocal called
@@ -391,7 +427,6 @@ async def use_context_manager():
         self.assertTrue(mock_instance.__aenter__.called)
         self.assertTrue(mock_instance.__aexit__.called)
         self.assertIsNot(mock_instance, result)
-        self.assertIsInstance(result, AsyncMock)
 
     def test_mock_customize_async_context_manager(self):
         instance = self.WithAsyncContextManager()

Misc/NEWS.d/next/Library/2019-09-11-14-45-30.bpo-38093.yQ6k7y.rst

@@ -0,0 +1,3 @@
+Fixes AsyncMock so MagicMocks on AsyncContextManagers do not have to
+manually set `__aenter__` and `__aexit__`. Also fixes AsyncMock so they
+don't crash when used with AsyncContextManager.