From e248174ca4a34d492bb375cf9bec14fa4205c3d6 Mon Sep 17 00:00:00 2001
From: Elvis Pranskevichus <elvis@magic.io>
Date: Wed, 26 Aug 2020 09:42:22 -0700
Subject: [PATCH] [3.7] bpo-32751: Wait for task cancel in asyncio.wait_for()
 when timeout <= 0 (GH-21895)

When I was fixing bpo-32751 back in GH-7216 I missed the case when
*timeout* is zero or negative.  This takes care of that.

Props to @aaliddell for noticing the inconsistency..
(cherry picked from commit c517fc712105c8e5930cb42baaebdbe37fc3e15f)

Co-authored-by: Elvis Pranskevichus <elvis@magic.io>
---
 Lib/asyncio/tasks.py                          |  9 ++++--
 Lib/test/test_asyncio/test_tasks.py           | 31 +++++++++++++++++++
 .../2020-08-15-15-50-12.bpo-32751.85je5X.rst  |  3 ++
 3 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2020-08-15-15-50-12.bpo-32751.85je5X.rst

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index c82b95e3beeaf2..5315f98300e10e 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -419,8 +419,13 @@ async def wait_for(fut, timeout, *, loop=None):
         if fut.done():
             return fut.result()
 
-        fut.cancel()
-        raise futures.TimeoutError()
+        await _cancel_and_wait(fut, loop=loop)
+        try:
+            fut.result()
+        except futures.CancelledError as exc:
+            raise futures.TimeoutError() from exc
+        else:
+            raise futures.TimeoutError()
 
     waiter = loop.create_future()
     timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index e0dbc95001d608..95e0188908c965 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -836,6 +836,9 @@ async def inner():
                 nonlocal task_done
                 try:
                     await asyncio.sleep(0.2, loop=loop)
+                except asyncio.CancelledError:
+                    await asyncio.sleep(0.1, loop=loop)
+                    raise
                 finally:
                     task_done = True
 
@@ -848,6 +851,34 @@ async def inner():
 
         loop.run_until_complete(foo())
 
+    def test_wait_for_waits_for_task_cancellation_w_timeout_0(self):
+        loop = asyncio.new_event_loop()
+        self.addCleanup(loop.close)
+
+        task_done = False
+
+        async def foo():
+            async def inner():
+                nonlocal task_done
+                try:
+                    await asyncio.sleep(10)
+                except asyncio.CancelledError:
+                    await asyncio.sleep(0.1)
+                    raise
+                finally:
+                    task_done = True
+
+            inner_task = self.new_task(loop, inner())
+            await asyncio.sleep(0.1)
+            await asyncio.wait_for(inner_task, timeout=0)
+
+        with self.assertRaises(asyncio.TimeoutError) as cm:
+            loop.run_until_complete(foo())
+
+        self.assertTrue(task_done)
+        chained = cm.exception.__context__
+        self.assertEqual(type(chained), asyncio.CancelledError)
+
     def test_wait_for_self_cancellation(self):
         loop = asyncio.new_event_loop()
         self.addCleanup(loop.close)
diff --git a/Misc/NEWS.d/next/Library/2020-08-15-15-50-12.bpo-32751.85je5X.rst b/Misc/NEWS.d/next/Library/2020-08-15-15-50-12.bpo-32751.85je5X.rst
new file mode 100644
index 00000000000000..c172ce5d9e9487
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-08-15-15-50-12.bpo-32751.85je5X.rst
@@ -0,0 +1,3 @@
+When cancelling the task due to a timeout, :meth:`asyncio.wait_for` will now
+wait until the cancellation is complete also in the case when *timeout* is
+<= 0, like it does with positive timeouts.