More tests for MapAsyncIterator

Author: Cito

README.md

@@ -13,7 +13,7 @@ a query language for APIs created by Facebook.
 
 The current version 1.0.0rc2 of GraphQL-core-next is up-to-date with GraphQL.js
 version 14.0.0rc2. All parts of the API are covered by an extensive test
-suite of currently 1561 unit tests.
+suite of currently 1562 unit tests.
 
 
 ## Documentation

graphql/subscription/map_async_iterator.py

@@ -19,49 +19,45 @@ def __init__(self, iterable: AsyncIterable, callback: Callable,
         self.iterator = iterable.__aiter__()
         self.callback = callback
         self.reject_callback = reject_callback
-        self.error = None
+        self.stop = False
 
     def __aiter__(self):
         return self
 
     async def __anext__(self):
-        if self.error is not None:
-            raise self.error
+        if self.stop:
+            raise StopAsyncIteration
         try:
             value = await self.iterator.__anext__()
         except Exception as error:
             if not self.reject_callback or isinstance(error, (
                     StopAsyncIteration, GeneratorExit)):
                 raise
-            if self.error is not None:
-                raise self.error
             result = self.reject_callback(error)
         else:
-            if self.error is not None:
-                raise self.error
             result = self.callback(value)
         if isawaitable(result):
             result = await result
-            if self.error is not None:
-                raise self.error
         return result
 
     async def athrow(self, type_, value=None, traceback=None):
-        if self.error:
+        if self.stop:
             return
         athrow = getattr(self.iterator, 'athrow', None)
         if athrow:
             await athrow(type_, value, traceback)
         else:
-            error = type_
-            if value is not None:
-                error = error(value)
-                if traceback is not None:
-                    error = error.with_traceback(traceback)
-            self.error = error
+            self.stop = True
+            if value is None:
+                if traceback is None:
+                    raise type_
+                value = type_()
+            if traceback is not None:
+                value = value.with_traceback(traceback)
+            raise value
 
     async def aclose(self):
-        if self.error:
+        if self.stop:
             return
         aclose = getattr(self.iterator, 'aclose', None)
         if aclose:
@@ -70,4 +66,4 @@ async def aclose(self):
             except RuntimeError:
                 pass
         else:
-            self.error = StopAsyncIteration
+            self.stop = True

tests/subscription/test_map_async_iterator.py

@@ -1,3 +1,5 @@
+import sys
+
 from pytest import mark, raises
 
 from graphql.subscription.map_async_iterator import MapAsyncIterator
@@ -171,3 +173,77 @@ async def source():
 
         with raises(StopAsyncIteration):
             await anext(doubles)
+
+    @mark.asyncio
+    async def can_use_simple_iterator_instead_of_generator():
+        async def source():
+            yield 1
+            yield 2
+            yield 3
+
+        class Source:
+            def __init__(self):
+                self.counter = 0
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                self.counter += 1
+                if self.counter > 3:
+                    raise StopAsyncIteration
+                return self.counter
+
+        for iterator in source, Source:
+            doubles = MapAsyncIterator(iterator(), lambda x: x + x)
+
+            await doubles.aclose()
+
+            with raises(StopAsyncIteration):
+                await anext(doubles)
+
+            doubles = MapAsyncIterator(iterator(), lambda x: x + x)
+
+            assert await anext(doubles) == 2
+            assert await anext(doubles) == 4
+            assert await anext(doubles) == 6
+
+            with raises(StopAsyncIteration):
+                await anext(doubles)
+
+            doubles = MapAsyncIterator(iterator(), lambda x: x + x)
+
+            assert await anext(doubles) == 2
+            assert await anext(doubles) == 4
+
+            # Throw error
+            with raises(RuntimeError) as exc_info:
+                await doubles.athrow(RuntimeError('ouch'))
+
+            assert str(exc_info.value) == 'ouch'
+
+            with raises(StopAsyncIteration):
+                await anext(doubles)
+            with raises(StopAsyncIteration):
+                await anext(doubles)
+
+            await doubles.athrow(RuntimeError('no more ouch'))
+
+            with raises(StopAsyncIteration):
+                await anext(doubles)
+
+            await doubles.aclose()
+
+            doubles = MapAsyncIterator(iterator(), lambda x: x + x)
+
+            assert await anext(doubles) == 2
+            assert await anext(doubles) == 4
+
+            try:
+                raise ValueError('bad')
+            except ValueError:
+                tb = sys.exc_info()[2]
+
+            # Throw error
+            with raises(ValueError):
+                await doubles.athrow(ValueError, None, tb)