GH-96704: Add {Task,Handle}.get_context(), use it in call_exception_handler() (#96756)

Co-authored-by: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>

1 files changed, 41 insertions, 0 deletions

diff --git a/Lib/test/test_asyncio/test_futures2.py b/Lib/test/test_asyncio/test_futures2.py
index 71279b69c79..9e7a5775a70 100644
--- a/Lib/test/test_asyncio/test_futures2.py
+++ b/Lib/test/test_asyncio/test_futures2.py
@@ -1,5 +1,6 @@
 # IsolatedAsyncioTestCase based tests
 import asyncio
+import contextvars
 import traceback
 import unittest
 from asyncio import tasks
@@ -27,6 +28,46 @@ class FutureTests:
             else:
                 self.fail('TypeError was not raised')
 
+    async def test_task_exc_handler_correct_context(self):
+        # see https://github.com/python/cpython/issues/96704
+        name = contextvars.ContextVar('name', default='foo')
+        exc_handler_called = False
+
+        def exc_handler(*args):
+            self.assertEqual(name.get(), 'bar')
+            nonlocal exc_handler_called
+            exc_handler_called = True
+
+        async def task():
+            name.set('bar')
+            1/0
+
+        loop = asyncio.get_running_loop()
+        loop.set_exception_handler(exc_handler)
+        self.cls(task())
+        await asyncio.sleep(0)
+        self.assertTrue(exc_handler_called)
+
+    async def test_handle_exc_handler_correct_context(self):
+        # see https://github.com/python/cpython/issues/96704
+        name = contextvars.ContextVar('name', default='foo')
+        exc_handler_called = False
+
+        def exc_handler(*args):
+            self.assertEqual(name.get(), 'bar')
+            nonlocal exc_handler_called
+            exc_handler_called = True
+
+        def callback():
+            name.set('bar')
+            1/0
+
+        loop = asyncio.get_running_loop()
+        loop.set_exception_handler(exc_handler)
+        loop.call_soon(callback)
+        await asyncio.sleep(0)
+        self.assertTrue(exc_handler_called)
+
 @unittest.skipUnless(hasattr(tasks, '_CTask'),
                        'requires the C _asyncio module')
 class CFutureTests(FutureTests, unittest.IsolatedAsyncioTestCase):