bpo-32751: Wait for task cancel in asyncio.wait_for() when timeout <= 0 (#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.

1 files changed, 31 insertions, 0 deletions

diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index f9db066..511961c 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1131,6 +1131,9 @@ class BaseTaskTests:
                 nonlocal task_done
                 try:
                     await asyncio.sleep(0.2)
+                except asyncio.CancelledError:
+                    await asyncio.sleep(_EPSILON)
+                    raise
                 finally:
                     task_done = True
 
@@ -1145,6 +1148,34 @@ class BaseTaskTests:
         chained = cm.exception.__context__
         self.assertEqual(type(chained), asyncio.CancelledError)
 
+    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(_EPSILON)
+                    raise
+                finally:
+                    task_done = True
+
+            inner_task = self.new_task(loop, inner())
+            await asyncio.sleep(_EPSILON)
+            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_reraises_exception_during_cancellation(self):
         loop = asyncio.new_event_loop()
         self.addCleanup(loop.close)