bpo-42130: Add test for swallowed cancellation

python · ods · May 12, 2021 · May 13, 2021 · May 13, 2021 · May 13, 2021
commit acb1522671591d978c5f6936f7d6c7c8ef5534d9
@@ -28,6 +28,7 @@ async def run(self):
 
         self.exited = True
 
+
 class AsyncioWaitForTest(unittest.TestCase):
 
     async def atest_asyncio_wait_for_cancelled(self):
@@ -56,6 +57,38 @@ async def atest_asyncio_wait_for_timeout(self):
     def test_asyncio_wait_for_timeout(self):
         asyncio.run(self.atest_asyncio_wait_for_timeout())
 
+    async def atest_asyncio_wait_for_hang(self, inner_steps, outer_steps):
+        # bpo-42130: wait_for can swallow cancellation causing task to hang
+        async def inner(steps):
+            for _ in range(steps):
+                await asyncio.sleep(0)
+
+        finished = False
+
+        async def wait_for_coro(steps):
+            nonlocal finished
+            await asyncio.wait_for(inner(steps), timeout=1)
+            await asyncio.sleep(0.1)
+            finished = True
+
+        task = asyncio.create_task(wait_for_coro(inner_steps))
+        for _ in range(outer_steps):
+            await asyncio.sleep(0)
+        assert not task.done()
+
+        task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await task
+        self.assertFalse(finished)
+
+    def test_asyncio_wait_for_hang(self):
+        # Test with different number of inner/outer steps to weaken dependence
+        # on implementation details
+        for inner_steps in range(3):
+            for outer_steps in range(3):
+                with self.subTest(inner_steps=inner_steps, outer_steps=outer_steps):
+                    asyncio.run(self.atest_asyncio_wait_for_hang(inner_steps, outer_steps))
+
 
 if __name__ == '__main__':
     unittest.main()