gh-129874: improve asyncio tests to use correct internal functions (#129887)

1 files changed, 26 insertions, 5 deletions

diff --git a/Lib/test/test_asyncio/test_free_threading.py b/Lib/test/test_asyncio/test_free_threading.py
index d0221d87062..199dbbdda5e 100644
--- a/Lib/test/test_asyncio/test_free_threading.py
+++ b/Lib/test/test_asyncio/test_free_threading.py
@@ -40,7 +40,7 @@ class TestFreeThreading:
                     self.assertEqual(task.get_loop(), loop)
                     self.assertFalse(task.done())
 
-                current = self.current_task()
+                current = asyncio.current_task()
                 self.assertEqual(current.get_loop(), loop)
                 self.assertSetEqual(all_tasks, tasks | {current})
                 future.set_result(None)
@@ -101,8 +101,12 @@ class TestFreeThreading:
         async def func():
             nonlocal task
             task = asyncio.current_task()
-
-        thread = Thread(target=lambda: asyncio.run(func()))
+        def runner():
+            with asyncio.Runner() as runner:
+                loop = runner.get_loop()
+                loop.set_task_factory(self.factory)
+                runner.run(func())
+        thread = Thread(target=runner)
         thread.start()
         thread.join()
         wr = weakref.ref(task)
@@ -164,7 +168,15 @@ class TestFreeThreading:
 
 class TestPyFreeThreading(TestFreeThreading, TestCase):
     all_tasks = staticmethod(asyncio.tasks._py_all_tasks)
-    current_task = staticmethod(asyncio.tasks._py_current_task)
+
+    def setUp(self):
+        self._old_current_task = asyncio.current_task
+        asyncio.current_task = asyncio.tasks.current_task = asyncio.tasks._py_current_task
+        return super().setUp()
+
+    def tearDown(self):
+        asyncio.current_task = asyncio.tasks.current_task = self._old_current_task
+        return super().tearDown()
 
     def factory(self, loop, coro, **kwargs):
         return asyncio.tasks._PyTask(coro, loop=loop, **kwargs)
@@ -173,7 +185,16 @@ class TestPyFreeThreading(TestFreeThreading, TestCase):
 @unittest.skipUnless(hasattr(asyncio.tasks, "_c_all_tasks"), "requires _asyncio")
 class TestCFreeThreading(TestFreeThreading, TestCase):
     all_tasks = staticmethod(getattr(asyncio.tasks, "_c_all_tasks", None))
-    current_task = staticmethod(getattr(asyncio.tasks, "_c_current_task", None))
+
+    def setUp(self):
+        self._old_current_task = asyncio.current_task
+        asyncio.current_task = asyncio.tasks.current_task = asyncio.tasks._c_current_task
+        return super().setUp()
+
+    def tearDown(self):
+        asyncio.current_task = asyncio.tasks.current_task = self._old_current_task
+        return super().tearDown()
+
 
     def factory(self, loop, coro, **kwargs):
         return asyncio.tasks._CTask(coro, loop=loop, **kwargs)