Fixes test_unittest changing env because of version issue.

python · lisroach · May 20, 2019 · May 14, 2018 · Jul 22, 2018 · Aug 16, 2018
commit a4d4dbceefd66b3485679d5b2dc13ef19b56ce49
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
@@ -1920,6 +1920,22 @@ def method(self, *args, **kw):
     '__index__': 1,
 }
 
+class AsyncIterator:
+    """
+    Wraps an iterator in an asynchronous iterator.
+    """
+    def __init__(self, iterator):
+        self.iterator = iterator
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self.iterator)
+        except StopIteration:
+            pass
+        raise StopAsyncIteration
 
 def _get_eq(self):
     def __eq__(other):
@@ -1950,10 +1966,26 @@ def __iter__():
         return iter(ret_val)
     return __iter__
 
+def _get_async_iter(mock):
+    def __aiter__():
+        return_value = mock.__aiter__._mock_return_value
+        if return_value is DEFAULT:
+            iterator = iter([])
+        else:
+            iterator = iter(return_value)
+
+        return AsyncIterator(iterator)
+
+    if asyncio.iscoroutinefunction(mock.__aiter__):
+        return asyncio.coroutine(__aiter__)
+
+    return __aiter__
+
 _side_effect_methods = {
     '__eq__': _get_eq,
     '__ne__': _get_ne,
     '__iter__': _get_iter,
+    '__aiter__': _get_async_iter
 }
 
 

diff --git a/Lib/unittest/test/testmock/testasync.py b/Lib/unittest/test/testmock/testasync.py
@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import sys
 import unittest
 
 from unittest.mock import call, CoroutineMock, patch, MagicMock
@@ -237,7 +238,7 @@ async def __aenter__(self, *args, **kwargs):
         async def __aexit__(self, *args, **kwargs):
             self.exited = True
 
-    class AsyncIterator:
+    class AsyncIterator(object):
         def __init__(self):
             self.iter_called = False
             self.next_called = False
@@ -254,6 +255,11 @@ async def __anext__(self):
 
             raise StopAsyncIteration
 
+    # Before 3.7 __aiter__ was a coroutine
+    class AsyncItertorDeprecated(AsyncIterator):
+        async def __aiter__(self):
+            return super().__aiter__()
+
     def test_mock_magic_methods_are_coroutine_mocks(self):
         mock_instance = CoroutineMock(spec=self.AsyncContextManager())
         self.assertIsInstance(mock_instance.__aenter__,
@@ -262,26 +268,27 @@ def test_mock_magic_methods_are_coroutine_mocks(self):
                               CoroutineMock)
 
     def test_mock_aiter_and_anext(self):
-        instance = self.AsyncIterator()
+        if sys.version_info < (3, 7):
+            instance = self.AsyncItertorDeprecated()
+        else:
+            instance = self.AsyncIterator()
         mock_instance = CoroutineMock(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):
+        if sys.version_info < (3, 7):
+            iterator = instance.__aiter__()
             iterator = run_coroutine(iterator)
 
-        mock_iterator = mock_instance.__aiter__()
-        if asyncio.iscoroutine(mock_iterator):
+            mock_iterator = mock_instance.__aiter__()
             mock_iterator = run_coroutine(mock_iterator)
 
-        self.assertEqual(asyncio.iscoroutine(iterator.__aiter__),
-                         asyncio.iscoroutine(mock_iterator.__aiter__))
-        self.assertEqual(asyncio.iscoroutine(iterator.__anext__),
-                         asyncio.iscoroutine(mock_iterator.__anext__))
+            self.assertEqual(asyncio.iscoroutine(iterator.__aiter__),
+                             asyncio.iscoroutine(mock_iterator.__aiter__))
+            self.assertEqual(asyncio.iscoroutine(iterator.__anext__),
+                             asyncio.iscoroutine(mock_iterator.__anext__))
 
 
 class CoroutineMockAssert(unittest.TestCase):