1 files changed, 129 insertions, 0 deletions
diff --git a/tests/extmod/asyncio_wait_for.py b/tests/extmod/asyncio_wait_for.py
new file mode 100644
index 0000000000..f973e46f4c
--- /dev/null
+++ b/tests/extmod/asyncio_wait_for.py
@@ -0,0 +1,129 @@
+# Test asyncio.wait_for
+
+try:
+    import asyncio
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+
+async def task(id, t):
+    print("task start", id)
+    await asyncio.sleep(t)
+    print("task end", id)
+    return id * 2
+
+
+async def task_catch():
+    print("task_catch start")
+    try:
+        await asyncio.sleep(0.2)
+    except asyncio.CancelledError:
+        print("ignore cancel")
+    print("task_catch done")
+
+
+async def task_raise():
+    print("task start")
+    raise ValueError
+
+
+async def task_cancel_other(t, other):
+    print("task_cancel_other start")
+    await asyncio.sleep(t)
+    print("task_cancel_other cancel")
+    other.cancel()
+
+
+async def task_wait_for_cancel(id, t, t_wait):
+    print("task_wait_for_cancel start")
+    try:
+        await asyncio.wait_for(task(id, t), t_wait)
+    except asyncio.CancelledError as er:
+        print("task_wait_for_cancel cancelled")
+        raise er
+
+
+async def task_wait_for_cancel_ignore(t_wait):
+    print("task_wait_for_cancel_ignore start")
+    try:
+        await asyncio.wait_for(task_catch(), t_wait)
+    except asyncio.CancelledError as er:
+        print("task_wait_for_cancel_ignore cancelled")
+        raise er
+
+
+async def main():
+    sep = "-" * 10
+
+    # When task finished before the timeout
+    print(await asyncio.wait_for(task(1, 0.01), 10))
+    print(sep)
+
+    # When timeout passes and task is cancelled
+    try:
+        print(await asyncio.wait_for(task(2, 10), 0.01))
+    except asyncio.TimeoutError:
+        print("timeout")
+    print(sep)
+
+    # When timeout passes and task is cancelled, but task ignores the cancellation request
+    try:
+        print(await asyncio.wait_for(task_catch(), 0.1))
+    except asyncio.TimeoutError:
+        print("TimeoutError")
+    print(sep)
+
+    # When task raises an exception
+    try:
+        print(await asyncio.wait_for(task_raise(), 1))
+    except ValueError:
+        print("ValueError")
+    print(sep)
+
+    # Timeout of None means wait forever
+    print(await asyncio.wait_for(task(3, 0.1), None))
+    print(sep)
+
+    # When task is cancelled by another task
+    t = asyncio.create_task(task(4, 10))
+    asyncio.create_task(task_cancel_other(0.01, t))
+    try:
+        print(await asyncio.wait_for(t, 1))
+    except asyncio.CancelledError as er:
+        print(repr(er))
+    print(sep)
+
+    # When wait_for gets cancelled
+    t = asyncio.create_task(task_wait_for_cancel(4, 1, 2))
+    await asyncio.sleep(0.01)
+    t.cancel()
+    await asyncio.sleep(0.01)
+    print(sep)
+
+    # When wait_for gets cancelled and awaited task ignores the cancellation request
+    t = asyncio.create_task(task_wait_for_cancel_ignore(2))
+    await asyncio.sleep(0.01)
+    t.cancel()
+    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")
+
+
+asyncio.run(main())