Add utility to wait for several events.

Indeed, trio has no equivalent of asyncio.wait.

Author: aaugustin

src/websockets/trio/utils.py

@@ -0,0 +1,42 @@
+import sys
+
+import trio
+
+
+if sys.version_info[:2] < (3, 11):  # pragma: no cover
+    from exceptiongroup import BaseExceptionGroup
+
+
+__all__ = ["race_events"]
+
+
+# Based on https://trio.readthedocs.io/en/stable/reference-core.html#custom-supervisors
+
+
+async def jockey(event: trio.Event, cancel_scope: trio.CancelScope) -> None:
+    await event.wait()
+    cancel_scope.cancel()
+
+
+async def race_events(*events: trio.Event) -> None:
+    """
+    Wait for any of the given events to be set.
+
+    Args:
+        *events: The events to wait for.
+
+    """
+    if not events:
+        raise ValueError("no events provided")
+
+    try:
+        async with trio.open_nursery() as nursery:
+            for event in events:
+                nursery.start_soon(jockey, event, nursery.cancel_scope)
+    except BaseExceptionGroup as exc:
+        try:
+            trio._util.raise_single_exception_from_group(exc)
+        except trio._util.MultipleExceptionError:  # pragma: no cover
+            raise AssertionError(
+                "race_events should be canceled; please file a bug report"
+            ) from exc

tests/trio/test_utils.py

@@ -0,0 +1,40 @@
+import trio.testing
+
+from websockets.trio.utils import *
+
+from .utils import IsolatedTrioTestCase
+
+
+class UtilsTests(IsolatedTrioTestCase):
+    async def test_race_events(self):
+        event1 = trio.Event()
+        event2 = trio.Event()
+        done = trio.Event()
+
+        async def waiter():
+            await race_events(event1, event2)
+            done.set()
+
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(waiter)
+            await trio.testing.wait_all_tasks_blocked()
+            self.assertFalse(done.is_set())
+
+            event1.set()
+            await trio.testing.wait_all_tasks_blocked()
+            self.assertTrue(done.is_set())
+
+    async def test_race_events_cancelled(self):
+        event1 = trio.Event()
+        event2 = trio.Event()
+
+        async def waiter():
+            with trio.move_on_after(0):
+                await race_events(event1, event2)
+
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(waiter)
+
+    async def test_race_events_no_events(self):
+        with self.assertRaises(ValueError):
+            await race_events()

tests/trio/utils.py

@@ -0,0 +1,62 @@
+import functools
+import inspect
+import sys
+import unittest
+
+import trio.testing
+
+
+if sys.version_info[:2] < (3, 11):  # pragma: no cover
+    from exceptiongroup import BaseExceptionGroup
+
+
+class IsolatedTrioTestCase(unittest.TestCase):
+    """
+    Wrap test coroutines with :func:`trio.testing.trio_test` automatically.
+
+    Create a nursery for each test, available in the :attr:`nursery` attribute.
+
+    :meth:`asyncSetUp` and :meth:`asyncTearDown` are supported, similar to
+    :class:`unittest.IsolatedAsyncioTestCase`, but ``addAsyncCleanup`` isn't.
+
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        for name in unittest.defaultTestLoader.getTestCaseNames(cls):
+            test = getattr(cls, name)
+            if getattr(test, "converted_to_trio", False):
+                return
+            assert inspect.iscoroutinefunction(test)
+            setattr(cls, name, cls.convert_to_trio(test))
+
+    @staticmethod
+    def convert_to_trio(test):
+        @trio.testing.trio_test
+        @functools.wraps(test)
+        async def new_test(self, *args, **kwargs):
+            try:
+                # Provide a nursery so it's easy to start tasks.
+                async with trio.open_nursery() as self.nursery:
+                    await self.asyncSetUp()
+                    try:
+                        return await test(self, *args, **kwargs)
+                    finally:
+                        await self.asyncTearDown()
+            except BaseExceptionGroup as exc:
+                # Unwrap exceptions like unittest.SkipTest. Multiple exceptions
+                # could occur is a test fails with multiple errors; this is OK;
+                # raise the original exception group in that case.
+                try:
+                    trio._util.raise_single_exception_from_group(exc)
+                except trio._util.MultipleExceptionError:  # pragma: no cover
+                    raise exc
+
+        new_test.converted_to_trio = True
+        return new_test
+
+    async def asyncSetUp(self):
+        pass
+
+    async def asyncTearDown(self):
+        pass