[3.8] bpo-38093: Correctly returns AsyncMock for async subclasses. (GH-15947) (GH-16299)

(cherry picked from commit 8b03f943c37e07fb2394acdcfacd066647f9b1fd)

Co-authored-by: Lisa Roach <lisaroach14@gmail.com>

3 files changed, 161 insertions, 63 deletions

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index 9fd5c3c..0c7545b 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -988,9 +988,13 @@ class NonCallableMock(Base):
         _type = type(self)
         if issubclass(_type, MagicMock) and _new_name in _async_method_magics:
             klass = AsyncMock
-        if issubclass(_type, AsyncMockMixin):
+        elif _new_name  in _sync_async_magics:
+            # Special case these ones b/c users will assume they are async,
+            # but they are actually sync (ie. __aiter__)
             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) :
@@ -1867,7 +1871,7 @@ _non_defaults = {
     '__reduce__', '__reduce_ex__', '__getinitargs__', '__getnewargs__',
     '__getstate__', '__setstate__', '__getformat__', '__setformat__',
     '__repr__', '__dir__', '__subclasses__', '__format__',
-    '__getnewargs_ex__', '__aenter__', '__aexit__', '__anext__', '__aiter__',
+    '__getnewargs_ex__',
 }
 
 
@@ -1886,10 +1890,12 @@ _magics = {
 
 # Magic methods used for async `with` statements
 _async_method_magics = {"__aenter__", "__aexit__", "__anext__"}
-# `__aiter__` is a plain function but used with async calls
-_async_magics = _async_method_magics | {"__aiter__"}
+# Magic methods that are only used with async calls but are synchronous functions themselves
+_sync_async_magics = {"__aiter__"}
+_async_magics = _async_method_magics | _sync_async_magics
 
-_all_magics = _magics | _non_defaults
+_all_sync_magics = _magics | _non_defaults
+_all_magics = _all_sync_magics | _async_magics
 
 _unsupported_magics = {
     '__getattr__', '__setattr__',
diff --git a/Lib/unittest/test/testmock/testasync.py b/Lib/unittest/test/testmock/testasync.py
index 4660bea..9b50b6d 100644
--- a/Lib/unittest/test/testmock/testasync.py
+++ b/Lib/unittest/test/testmock/testasync.py
@@ -377,10 +377,6 @@ class AsyncArguments(unittest.TestCase):
 
 class AsyncContextManagerTest(unittest.TestCase):
     class WithAsyncContextManager:
-        def __init__(self):
-            self.entered = False
-            self.exited = False
-
         async def __aenter__(self, *args, **kwargs):
             self.entered = True
             return self
@@ -388,32 +384,81 @@ class AsyncContextManagerTest(unittest.TestCase):
         async def __aexit__(self, *args, **kwargs):
             self.exited = True
 
-    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
+        def __init__(self):
+            self.session = None
+
+        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_magicmock_has_async_magic_methods(self):
+        cm = MagicMock(name='magic_cm')
+        self.assertTrue(hasattr(cm, "__aenter__"))
+        self.assertTrue(hasattr(cm, "__aexit__"))
+
+    def test_magic_methods_are_async_functions(self):
+        cm = MagicMock(name='magic_cm')
+        self.assertIsInstance(cm.__aenter__, AsyncMock)
+        self.assertIsInstance(cm.__aexit__, AsyncMock)
+        # AsyncMocks are also coroutine functions
+        self.assertTrue(asyncio.iscoroutinefunction(cm.__aenter__))
+        self.assertTrue(asyncio.iscoroutinefunction(cm.__aexit__))
+
+    def test_set_return_value_of_aenter(self):
+        def inner_test(mock_type):
+            pc = self.ProductionCode()
+            pc.session = MagicMock(name='sessionmock')
+            cm = mock_type(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})
+
+        for mock_type in [AsyncMock, MagicMock]:
+            with self.subTest(f"test set return value of aenter with {mock_type}"):
+                inner_test(mock_type)
 
     def test_mock_supports_async_context_manager(self):
-        called = False
-        instance = self.WithAsyncContextManager()
-        mock_instance = MagicMock(instance)
+        def inner_test(mock_type):
+            called = False
+            cm = self.WithAsyncContextManager()
+            cm_mock = mock_type(cm)
+
+            async def use_context_manager():
+                nonlocal called
+                async with cm_mock as result:
+                    called = True
+                return result
 
-        async def use_context_manager():
-            nonlocal called
-            async with mock_instance as result:
-                called = True
-            return result
-
-        result = asyncio.run(use_context_manager())
-        self.assertFalse(instance.entered)
-        self.assertFalse(instance.exited)
-        self.assertTrue(called)
-        self.assertTrue(mock_instance.entered)
-        self.assertTrue(mock_instance.exited)
-        self.assertTrue(mock_instance.__aenter__.called)
-        self.assertTrue(mock_instance.__aexit__.called)
-        self.assertIsNot(mock_instance, result)
-        self.assertIsInstance(result, AsyncMock)
+            cm_result = asyncio.run(use_context_manager())
+            self.assertTrue(called)
+            self.assertTrue(cm_mock.__aenter__.called)
+            self.assertTrue(cm_mock.__aexit__.called)
+            cm_mock.__aenter__.assert_awaited()
+            cm_mock.__aexit__.assert_awaited()
+            # We mock __aenter__ so it does not return self
+            self.assertIsNot(cm_mock, cm_result)
+
+        for mock_type in [AsyncMock, MagicMock]:
+            with self.subTest(f"test context manager magics with {mock_type}"):
+                inner_test(mock_type)
 
     def test_mock_customize_async_context_manager(self):
         instance = self.WithAsyncContextManager()
@@ -481,27 +526,30 @@ class AsyncIteratorTest(unittest.TestCase):
 
             raise StopAsyncIteration
 
-    def test_mock_aiter_and_anext(self):
-        instance = self.WithAsyncIterator()
-        mock_instance = MagicMock(instance)
-
-        self.assertEqual(asyncio.iscoroutine(instance.__aiter__),
-                         asyncio.iscoroutine(mock_instance.__aiter__))
-        self.assertEqual(asyncio.iscoroutine(instance.__anext__),
-                         asyncio.iscoroutine(mock_instance.__anext__))
-
-        iterator = instance.__aiter__()
-        if asyncio.iscoroutine(iterator):
-            iterator = asyncio.run(iterator)
-
-        mock_iterator = mock_instance.__aiter__()
-        if asyncio.iscoroutine(mock_iterator):
-            mock_iterator = asyncio.run(mock_iterator)
+    def test_aiter_set_return_value(self):
+        mock_iter = AsyncMock(name="tester")
+        mock_iter.__aiter__.return_value = [1, 2, 3]
+        async def main():
+            return [i async for i in mock_iter]
+        result = asyncio.run(main())
+        self.assertEqual(result, [1, 2, 3])
+
+    def test_mock_aiter_and_anext_asyncmock(self):
+        def inner_test(mock_type):
+            instance = self.WithAsyncIterator()
+            mock_instance = mock_type(instance)
+            # Check that the mock and the real thing bahave the same
+            # __aiter__ is not actually async, so not a coroutinefunction
+            self.assertFalse(asyncio.iscoroutinefunction(instance.__aiter__))
+            self.assertFalse(asyncio.iscoroutinefunction(mock_instance.__aiter__))
+            # __anext__ is async
+            self.assertTrue(asyncio.iscoroutinefunction(instance.__anext__))
+            self.assertTrue(asyncio.iscoroutinefunction(mock_instance.__anext__))
+
+        for mock_type in [AsyncMock, MagicMock]:
+            with self.subTest(f"test aiter and anext corourtine with {mock_type}"):
+                inner_test(mock_type)
 
-        self.assertEqual(asyncio.iscoroutine(iterator.__aiter__),
-                         asyncio.iscoroutine(mock_iterator.__aiter__))
-        self.assertEqual(asyncio.iscoroutine(iterator.__anext__),
-                         asyncio.iscoroutine(mock_iterator.__anext__))
 
     def test_mock_async_for(self):
         async def iterate(iterator):
@@ -512,19 +560,30 @@ class AsyncIteratorTest(unittest.TestCase):
             return accumulator
 
         expected = ["FOO", "BAR", "BAZ"]
-        with self.subTest("iterate through default value"):
-            mock_instance = MagicMock(self.WithAsyncIterator())
-            self.assertEqual([], asyncio.run(iterate(mock_instance)))
+        def test_default(mock_type):
+            mock_instance = mock_type(self.WithAsyncIterator())
+            self.assertEqual(asyncio.run(iterate(mock_instance)), [])
+
 
-        with self.subTest("iterate through set return_value"):
-            mock_instance = MagicMock(self.WithAsyncIterator())
+        def test_set_return_value(mock_type):
+            mock_instance = mock_type(self.WithAsyncIterator())
             mock_instance.__aiter__.return_value = expected[:]
-            self.assertEqual(expected, asyncio.run(iterate(mock_instance)))
+            self.assertEqual(asyncio.run(iterate(mock_instance)), expected)
 
-        with self.subTest("iterate through set return_value iterator"):
-            mock_instance = MagicMock(self.WithAsyncIterator())
+        def test_set_return_value_iter(mock_type):
+            mock_instance = mock_type(self.WithAsyncIterator())
             mock_instance.__aiter__.return_value = iter(expected[:])
-            self.assertEqual(expected, asyncio.run(iterate(mock_instance)))
+            self.assertEqual(asyncio.run(iterate(mock_instance)), expected)
+
+        for mock_type in [AsyncMock, MagicMock]:
+            with self.subTest(f"default value with {mock_type}"):
+                test_default(mock_type)
+
+            with self.subTest(f"set return_value with {mock_type}"):
+                test_set_return_value(mock_type)
+
+            with self.subTest(f"set return_value iterator with {mock_type}"):
+                test_set_return_value_iter(mock_type)
 
 
 class AsyncMockAssert(unittest.TestCase):
diff --git a/Lib/unittest/test/testmock/testmagicmethods.py b/Lib/unittest/test/testmock/testmagicmethods.py
index 130a339..57f85e9 100644
--- a/Lib/unittest/test/testmock/testmagicmethods.py
+++ b/Lib/unittest/test/testmock/testmagicmethods.py
@@ -1,8 +1,9 @@
+import asyncio
 import math
 import unittest
 import os
 import sys
-from unittest.mock import Mock, MagicMock, _magics
+from unittest.mock import AsyncMock, Mock, MagicMock, _magics
 
 
 
@@ -271,6 +272,34 @@ class TestMockingMagicMethods(unittest.TestCase):
         self.assertEqual(mock != mock, False)
 
 
+    # This should be fixed with issue38163
+    @unittest.expectedFailure
+    def test_asyncmock_defaults(self):
+        mock = AsyncMock()
+        self.assertEqual(int(mock), 1)
+        self.assertEqual(complex(mock), 1j)
+        self.assertEqual(float(mock), 1.0)
+        self.assertNotIn(object(), mock)
+        self.assertEqual(len(mock), 0)
+        self.assertEqual(list(mock), [])
+        self.assertEqual(hash(mock), object.__hash__(mock))
+        self.assertEqual(str(mock), object.__str__(mock))
+        self.assertTrue(bool(mock))
+        self.assertEqual(round(mock), mock.__round__())
+        self.assertEqual(math.trunc(mock), mock.__trunc__())
+        self.assertEqual(math.floor(mock), mock.__floor__())
+        self.assertEqual(math.ceil(mock), mock.__ceil__())
+        self.assertTrue(asyncio.iscoroutinefunction(mock.__aexit__))
+        self.assertTrue(asyncio.iscoroutinefunction(mock.__aenter__))
+        self.assertIsInstance(mock.__aenter__, AsyncMock)
+        self.assertIsInstance(mock.__aexit__, AsyncMock)
+
+        # in Python 3 oct and hex use __index__
+        # so these tests are for __index__ in py3k
+        self.assertEqual(oct(mock), '0o1')
+        self.assertEqual(hex(mock), '0x1')
+        # how to test __sizeof__ ?
+
     def test_magicmock_defaults(self):
         mock = MagicMock()
         self.assertEqual(int(mock), 1)
@@ -286,6 +315,10 @@ class TestMockingMagicMethods(unittest.TestCase):
         self.assertEqual(math.trunc(mock), mock.__trunc__())
         self.assertEqual(math.floor(mock), mock.__floor__())
         self.assertEqual(math.ceil(mock), mock.__ceil__())
+        self.assertTrue(asyncio.iscoroutinefunction(mock.__aexit__))
+        self.assertTrue(asyncio.iscoroutinefunction(mock.__aenter__))
+        self.assertIsInstance(mock.__aenter__, AsyncMock)
+        self.assertIsInstance(mock.__aexit__, AsyncMock)
 
         # in Python 3 oct and hex use __index__
         # so these tests are for __index__ in py3k