summaryrefslogtreecommitdiffstatshomepage
path: root/tests/extmod/asyncio_wait_for.py
blob: f973e46f4cf91d4c5d538ae0d74826be04c3642e (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
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())