fix coverage

Author: rmorshea

src/py/reactpy/reactpy/core/hooks.py

@@ -3,7 +3,7 @@
 import inspect
 import sys
 import warnings
-from asyncio import CancelledError, Event, create_task
+from asyncio import FIRST_COMPLETED, CancelledError, Event, create_task, wait
 from collections.abc import Awaitable, Coroutine, Sequence
 from logging import getLogger
 from types import FunctionType
@@ -240,20 +240,36 @@ async def __aexit__(self, exc_type: type[BaseException], *exc: Any) -> Any:
             # propagate non-cancellation exceptions
             return None
 
+        if not self._maybe_uncancel_task():
+            return None
+
+        wait_for_stop = create_task(self._stop.wait())
         try:
-            await self._stop.wait()
+            await wait([wait_for_stop, self.task], return_when=FIRST_COMPLETED)
         except CancelledError:
-            if self.task.cancelling() > self._cancel_count:
-                # Task has been cancelled by something else - propagate it
-                return None
-            self.task.uncancel()
+            if not self._maybe_uncancel_task():
+                raise
 
         return True
 
     def _cancel_task(self) -> None:
         self.task.cancel()
         self._cancel_count += 1
 
+    def _maybe_uncancel_task(self) -> bool:
+        """Return if task was uncancelled
+
+        If task was not cancelled by this effect then returns. Otherwise, if the task
+        was cancelled at all, uncancell it and return True
+        """
+        if self.task.cancelling() > self._cancel_count:
+            # Task has been cancelled by something else - propagate it
+            return False
+        elif self._cancel_count:
+            for _ in range(self._cancel_count):
+                self.task.uncancel()
+        return True
+
 
 def _cast_async_effect(function: Callable[..., Any]) -> _AsyncEffectFunc:
     if inspect.iscoroutinefunction(function):

src/py/reactpy/tests/test_core/test_hooks.py

@@ -587,27 +587,46 @@ async def effect_func(e):
     await did_cleanup.wait()
 
 
-async def test_effect_external_cancellation_is_propagated():
+async def test_effect_internal_cancellation_is_propagated():
     did_start = WaitForEvent()
     did_cleanup = Ref(False)
+    never_happens = asyncio.Event()
 
     async def effect_func(e):
         async with e:
             did_start.set()
-            # We don't use e.task or asyncio.current_task() because there seems to
-            # be some platform dependence on whether the task is cancelled properly
-            outer_task.cancel()
-            # Allow cancellation to propagate
-            await asyncio.sleep(0)
+            asyncio.current_task().cancel()
+            await never_happens.wait()
         did_cleanup.current = True
 
     async def main():
         effect = Effect(effect_func)
         await did_start.wait()
-        await effect.stop()
+        await effect.task
+
+    with pytest.raises(asyncio.CancelledError):
+        await main()
+
+    assert not did_cleanup.current
+
+
+async def test_effect_external_cancellation_is_propagated():
+    did_start = WaitForEvent()
+    did_cleanup = Ref(False)
+
+    async def effect_func(e):
+        async with e:
+            did_start.set()
+        did_cleanup.current = True
+
+    async def main():
+        effect = Effect(effect_func)
+        await effect.task
 
     with pytest.raises(asyncio.CancelledError):
         outer_task = asyncio.create_task(main())
+        await did_start.wait()
+        outer_task.cancel()
         await outer_task
 
     assert not did_cleanup.current