1# Test asyncio.wait_for
2
3try:
4    import uasyncio as asyncio
5except ImportError:
6    try:
7        import asyncio
8    except ImportError:
9        print("SKIP")
10        raise SystemExit
11
12
13async def task(id, t):
14    print("task start", id)
15    await asyncio.sleep(t)
16    print("task end", id)
17    return id * 2
18
19
20async def task_catch():
21    print("task_catch start")
22    try:
23        await asyncio.sleep(0.2)
24    except asyncio.CancelledError:
25        print("ignore cancel")
26    print("task_catch done")
27
28
29async def task_raise():
30    print("task start")
31    raise ValueError
32
33
34async def task_cancel_other(t, other):
35    print("task_cancel_other start")
36    await asyncio.sleep(t)
37    print("task_cancel_other cancel")
38    other.cancel()
39
40
41async def task_wait_for_cancel(id, t, t_wait):
42    print("task_wait_for_cancel start")
43    try:
44        await asyncio.wait_for(task(id, t), t_wait)
45    except asyncio.CancelledError as er:
46        print("task_wait_for_cancel cancelled")
47        raise er
48
49
50async def task_wait_for_cancel_ignore(t_wait):
51    print("task_wait_for_cancel_ignore start")
52    try:
53        await asyncio.wait_for(task_catch(), t_wait)
54    except asyncio.CancelledError as er:
55        print("task_wait_for_cancel_ignore cancelled")
56        raise er
57
58
59async def main():
60    sep = "-" * 10
61
62    # When task finished before the timeout
63    print(await asyncio.wait_for(task(1, 0.01), 10))
64    print(sep)
65
66    # When timeout passes and task is cancelled
67    try:
68        print(await asyncio.wait_for(task(2, 10), 0.01))
69    except asyncio.TimeoutError:
70        print("timeout")
71    print(sep)
72
73    # When timeout passes and task is cancelled, but task ignores the cancellation request
74    try:
75        print(await asyncio.wait_for(task_catch(), 0.1))
76    except asyncio.TimeoutError:
77        print("TimeoutError")
78    print(sep)
79
80    # When task raises an exception
81    try:
82        print(await asyncio.wait_for(task_raise(), 1))
83    except ValueError:
84        print("ValueError")
85    print(sep)
86
87    # Timeout of None means wait forever
88    print(await asyncio.wait_for(task(3, 0.1), None))
89    print(sep)
90
91    # When task is cancelled by another task
92    t = asyncio.create_task(task(4, 10))
93    asyncio.create_task(task_cancel_other(0.01, t))
94    try:
95        print(await asyncio.wait_for(t, 1))
96    except asyncio.CancelledError as er:
97        print(repr(er))
98    print(sep)
99
100    # When wait_for gets cancelled
101    t = asyncio.create_task(task_wait_for_cancel(4, 1, 2))
102    await asyncio.sleep(0.01)
103    t.cancel()
104    await asyncio.sleep(0.01)
105    print(sep)
106
107    # When wait_for gets cancelled and awaited task ignores the cancellation request
108    t = asyncio.create_task(task_wait_for_cancel_ignore(2))
109    await asyncio.sleep(0.01)
110    t.cancel()
111    await asyncio.sleep(0.01)
112    print(sep)
113
114    print("finish")
115
116
117asyncio.run(main())
118