Merge pull request #27 from syrusakbary/await-blocked-by-context-fix

Fix await blocked by context exit

promise/promise.py

@@ -140,7 +140,12 @@ def future(self):
         return self._future
 
     def __iter__(self):
-        return iterate_promise(self)
+        if self._trace:
+            # If we wait, we drain the queue of the
+            # callbacks waiting on the context exit
+            # so we avoid a blocking state
+            self._trace.drain_queue()
+        return iterate_promise(self._target())
 
     __await__ = __iter__
 
@@ -643,6 +648,7 @@ def _try_convert_to_promise(cls, obj, context=None):
 
         if iscoroutine(obj):
             obj = ensure_future(obj)
+            _type = obj.__class__
 
         if is_future_like(_type):
             def executor(resolve, reject):

tests/conftest.py

@@ -5,3 +5,4 @@
     collect_ignore.append('test_awaitable.py')
 if version_info[:2] < (3, 5):
     collect_ignore.append('test_awaitable_35.py')
+    collect_ignore.append('test_dataloader_awaitable_35.py')

tests/test_awaitable.py

@@ -19,3 +19,14 @@ def resolve_or_reject(resolve, reject):
 
     p = Promise(resolve_or_reject)
     assert p.get() is True
+
+
+@mark.asyncio
+@coroutine
+def test_promise_coroutine():
+    @coroutine
+    def my_coro():
+        yield from Promise.resolve(True)
+
+    promise = Promise.resolve(my_coro())
+    assert isinstance(promise, Promise)

tests/test_awaitable_35.py

@@ -1,5 +1,6 @@
 from asyncio import sleep, Future, wait, FIRST_COMPLETED
 from pytest import mark
+from promise.context import Context
 from promise import Promise, is_thenable
 
 
@@ -31,3 +32,14 @@ async def test_promisify_future():
     future = Future()
     future.set_result(True)
     assert await Promise.resolve(future)
+
+
+@mark.asyncio
+async def test_await_in_context():
+    async def inner():
+        with Context():
+            promise = Promise.resolve(True).then(lambda x: x)
+            return await promise
+
+    result = await inner()
+    assert result == True

tests/test_dataloader_awaitable_35.py

@@ -0,0 +1,73 @@
+from pytest import mark
+from promise import Promise
+from promise.dataloader import DataLoader
+
+
+def id_loader(**options):
+    load_calls = []
+
+    resolve = options.pop('resolve', Promise.resolve)
+
+    def fn(keys):
+        load_calls.append(keys)
+        return resolve(keys)
+
+    identity_loader = DataLoader(fn, **options)
+    return identity_loader, load_calls
+
+
+@mark.asyncio
+async def test_await_dataloader():
+    identity_loader, load_calls = id_loader()
+    async def load_multiple(identity_loader):
+        one = identity_loader.load('load1')
+        two = identity_loader.load('load2')
+        return await Promise.all([one, two])
+
+    result = await load_multiple(identity_loader)
+    assert result == ['load1', 'load2']
+    assert load_calls == [['load1'], ['load2']]
+
+
+@mark.asyncio
+async def test_await_dataloader_safe_promise():
+    identity_loader, load_calls = id_loader()
+
+    @Promise.safe
+    async def load_multiple(identity_loader):
+        one = identity_loader.load('load1')
+        two = identity_loader.load('load2')
+        return await Promise.all([one, two])
+
+    result = await load_multiple(identity_loader)
+    assert result == ['load1', 'load2']
+    assert load_calls == [['load1'], ['load2']]
+
+
+@mark.asyncio
+async def test_await_dataloader_individual():
+    identity_loader, load_calls = id_loader()
+
+    async def load_one_then_two(identity_loader):
+        one = await identity_loader.load('load1')
+        two = await identity_loader.load('load2')
+        return [one, two]
+
+    result = await load_one_then_two(identity_loader)
+    assert result == ['load1', 'load2']
+    assert load_calls == [['load1'], ['load2']]
+
+
+@mark.asyncio
+async def test_await_dataloader_individual_safe_promise():
+    identity_loader, load_calls = id_loader()
+
+    @Promise.safe
+    async def load_one_then_two(identity_loader):
+        one = await identity_loader.load('load1')
+        two = await identity_loader.load('load2')
+        return [one, two]
+
+    result = await load_one_then_two(identity_loader)
+    assert result == ['load1', 'load2']
+    assert load_calls == [['load1'], ['load2']]