gh-124309: Modernize the `staggered_race` implementation to support eager task factories (#124390)

Co-authored-by: Thomas Grainger <tagrain@gmail.com>
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Co-authored-by: Carol Willing <carolcode@willingconsulting.com>
Co-authored-by: Kumar Aditya <kumaraditya@python.org>

2 files changed, 80 insertions, 4 deletions

diff --git a/Lib/test/test_asyncio/test_eager_task_factory.py b/Lib/test/test_asyncio/test_eager_task_factory.py
index 0777f39..1579ad1 100644
--- a/Lib/test/test_asyncio/test_eager_task_factory.py
+++ b/Lib/test/test_asyncio/test_eager_task_factory.py
@@ -213,6 +213,53 @@ class EagerTaskFactoryLoopTests:
 
         self.run_coro(run())
 
+    def test_staggered_race_with_eager_tasks(self):
+        # See https://github.com/python/cpython/issues/124309
+
+        async def fail():
+            await asyncio.sleep(0)
+            raise ValueError("no good")
+
+        async def run():
+            winner, index, excs = await asyncio.staggered.staggered_race(
+                [
+                    lambda: asyncio.sleep(2, result="sleep2"),
+                    lambda: asyncio.sleep(1, result="sleep1"),
+                    lambda: fail()
+                ],
+                delay=0.25
+            )
+            self.assertEqual(winner, 'sleep1')
+            self.assertEqual(index, 1)
+            self.assertIsNone(excs[index])
+            self.assertIsInstance(excs[0], asyncio.CancelledError)
+            self.assertIsInstance(excs[2], ValueError)
+
+        self.run_coro(run())
+
+    def test_staggered_race_with_eager_tasks_no_delay(self):
+        # See https://github.com/python/cpython/issues/124309
+        async def fail():
+            raise ValueError("no good")
+
+        async def run():
+            winner, index, excs = await asyncio.staggered.staggered_race(
+                [
+                    lambda: fail(),
+                    lambda: asyncio.sleep(1, result="sleep1"),
+                    lambda: asyncio.sleep(0, result="sleep0"),
+                ],
+                delay=None
+            )
+            self.assertEqual(winner, 'sleep1')
+            self.assertEqual(index, 1)
+            self.assertIsNone(excs[index])
+            self.assertIsInstance(excs[0], ValueError)
+            self.assertEqual(len(excs), 2)
+
+        self.run_coro(run())
+
+
 
 class PyEagerTaskFactoryLoopTests(EagerTaskFactoryLoopTests, test_utils.TestCase):
     Task = tasks._PyTask
diff --git a/Lib/test/test_asyncio/test_staggered.py b/Lib/test/test_asyncio/test_staggered.py
index e6e32f7..21a39b3 100644
--- a/Lib/test/test_asyncio/test_staggered.py
+++ b/Lib/test/test_asyncio/test_staggered.py
@@ -82,16 +82,45 @@ class StaggeredTests(unittest.IsolatedAsyncioTestCase):
         async def coro(index):
             raise ValueError(index)
 
+        for delay in [None, 0, 0.1, 1]:
+            with self.subTest(delay=delay):
+                winner, index, excs = await staggered_race(
+                    [
+                        lambda: coro(0),
+                        lambda: coro(1),
+                    ],
+                    delay=delay,
+                )
+
+                self.assertIs(winner, None)
+                self.assertIs(index, None)
+                self.assertEqual(len(excs), 2)
+                self.assertIsInstance(excs[0], ValueError)
+                self.assertIsInstance(excs[1], ValueError)
+
+    async def test_long_delay_early_failure(self):
+        async def coro(index):
+            await asyncio.sleep(0)  # Dummy coroutine for the 1 case
+            if index == 0:
+                await asyncio.sleep(0.1)  # Dummy coroutine
+                raise ValueError(index)
+
+            return f'Res: {index}'
+
         winner, index, excs = await staggered_race(
             [
                 lambda: coro(0),
                 lambda: coro(1),
             ],
-            delay=None,
+            delay=10,
         )
 
-        self.assertIs(winner, None)
-        self.assertIs(index, None)
+        self.assertEqual(winner, 'Res: 1')
+        self.assertEqual(index, 1)
         self.assertEqual(len(excs), 2)
         self.assertIsInstance(excs[0], ValueError)
-        self.assertIsInstance(excs[1], ValueError)
+        self.assertIsNone(excs[1])
+
+
+if __name__ == "__main__":
+    unittest.main()