add more thread safety tests

Author: kumaraditya303

Lib/test/test_asyncio/test_free_threading.py

@@ -7,6 +7,10 @@
 
 threading_helper.requires_working_threading(module=True)
 
+class MyException(Exception):
+    pass
+
+
 def tearDownModule():
     asyncio._set_event_loop_policy(None)
 
@@ -53,6 +57,88 @@ def runner():
         with threading_helper.start_threads(threads):
             pass
 
+    def test_run_coroutine_threadsafe(self) -> None:
+        results = []
+
+        def in_thread(loop: asyncio.AbstractEventLoop):
+            coro = asyncio.sleep(0.1, result=42)
+            fut = asyncio.run_coroutine_threadsafe(coro, loop)
+            result = fut.result()
+            self.assertEqual(result, 42)
+            results.append(result)
+
+        async def main():
+            loop = asyncio.get_running_loop()
+            async with asyncio.TaskGroup() as tg:
+                for _ in range(10):
+                    tg.create_task(asyncio.to_thread(in_thread, loop))
+            self.assertEqual(results, [42] * 10)
+
+        with asyncio.Runner() as r:
+            loop = r.get_loop()
+            loop.set_task_factory(self.factory)
+            r.run(main())
+
+    def test_run_coroutine_threadsafe_exception_caught(self) -> None:
+        exc = MyException("test")
+
+        async def coro():
+            await asyncio.sleep(0.1)
+            raise exc
+
+        def in_thread(loop: asyncio.AbstractEventLoop):
+            fut = asyncio.run_coroutine_threadsafe(coro(), loop)
+            self.assertEqual(fut.exception(), exc)
+            return exc
+
+        async def main():
+            loop = asyncio.get_running_loop()
+            tasks = []
+            async with asyncio.TaskGroup() as tg:
+                for _ in range(10):
+                    task = tg.create_task(asyncio.to_thread(in_thread, loop))
+                    tasks.append(task)
+            for task in tasks:
+                self.assertEqual(await task, exc)
+
+        with asyncio.Runner() as r:
+            loop = r.get_loop()
+            loop.set_task_factory(self.factory)
+            r.run(main())
+
+    def test_run_coroutine_threadsafe_exception_uncaught(self) -> None:
+        async def coro():
+            await asyncio.sleep(1)
+            raise MyException("test")
+
+        def in_thread(loop: asyncio.AbstractEventLoop):
+            fut = asyncio.run_coroutine_threadsafe(coro(), loop)
+            return fut.result()
+
+        async def main():
+            loop = asyncio.get_running_loop()
+            tasks = []
+            try:
+                async with asyncio.TaskGroup() as tg:
+                    for _ in range(10):
+                        task = tg.create_task(asyncio.to_thread(in_thread, loop))
+                        tasks.append(task)
+            except ExceptionGroup:
+                for task in tasks:
+                    try:
+                        await task
+                    except (MyException, asyncio.CancelledError):
+                        pass
+                    else:
+                        self.fail("Task should have raised an exception")
+            else:
+                self.fail("TaskGroup should have raised an exception")
+
+        with asyncio.Runner() as r:
+            loop = r.get_loop()
+            loop.set_task_factory(self.factory)
+            r.run(main())
+
 
 class TestPyFreeThreading(TestFreeThreading, TestCase):
     all_tasks = staticmethod(asyncio.tasks._py_all_tasks)