extmod/uasyncio: Fix edge case for cancellation of wait_for.

This fixes the cases where the task being waited on finishes just before or
just after the wait_for itself is cancelled.

Fixes issue #8717.

Signed-off-by: Damien George <damien@micropython.org>

2 files changed, 42 insertions, 0 deletions

diff --git a/tests/extmod/uasyncio_wait_for.py b/tests/extmod/uasyncio_wait_for.py
index 9612d1620..c636c7dd7 100644
--- a/tests/extmod/uasyncio_wait_for.py
+++ b/tests/extmod/uasyncio_wait_for.py
@@ -111,6 +111,21 @@ async def main():
     await asyncio.sleep(0.01)
     print(sep)
 
+    # When wait_for gets cancelled and the task it's waiting on finishes around the
+    # same time as the cancellation of the wait_for
+    for num_sleep in range(1, 5):
+        t = asyncio.create_task(task_wait_for_cancel(4 + num_sleep, 0, 2))
+        for _ in range(num_sleep):
+            await asyncio.sleep(0)
+        assert not t.done()
+        print("cancel wait_for")
+        t.cancel()
+        try:
+            await t
+        except asyncio.CancelledError as er:
+            print(repr(er))
+        print(sep)
+
     print("finish")
 
 
diff --git a/tests/extmod/uasyncio_wait_for.py.exp b/tests/extmod/uasyncio_wait_for.py.exp
index a4201d31f..1bbe3d065 100644
--- a/tests/extmod/uasyncio_wait_for.py.exp
+++ b/tests/extmod/uasyncio_wait_for.py.exp
@@ -32,4 +32,31 @@ task_wait_for_cancel_ignore cancelled
 ignore cancel
 task_catch done
 ----------
+task_wait_for_cancel start
+cancel wait_for
+task start 5
+task_wait_for_cancel cancelled
+CancelledError()
+----------
+task_wait_for_cancel start
+task start 6
+cancel wait_for
+task end 6
+task_wait_for_cancel cancelled
+CancelledError()
+----------
+task_wait_for_cancel start
+task start 7
+task end 7
+cancel wait_for
+task_wait_for_cancel cancelled
+CancelledError()
+----------
+task_wait_for_cancel start
+task start 8
+task end 8
+cancel wait_for
+task_wait_for_cancel cancelled
+CancelledError()
+----------
 finish