• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import asyncio
2import unittest
3import time
4
5def tearDownModule():
6    asyncio.set_event_loop_policy(None)
7
8
9class SlowTask:
10    """ Task will run for this defined time, ignoring cancel requests """
11    TASK_TIMEOUT = 0.2
12
13    def __init__(self):
14        self.exited = False
15
16    async def run(self):
17        exitat = time.monotonic() + self.TASK_TIMEOUT
18
19        while True:
20            tosleep = exitat - time.monotonic()
21            if tosleep <= 0:
22                break
23
24            try:
25                await asyncio.sleep(tosleep)
26            except asyncio.CancelledError:
27                pass
28
29        self.exited = True
30
31class AsyncioWaitForTest(unittest.TestCase):
32
33    async def atest_asyncio_wait_for_cancelled(self):
34        t  = SlowTask()
35
36        waitfortask = asyncio.create_task(asyncio.wait_for(t.run(), t.TASK_TIMEOUT * 2))
37        await asyncio.sleep(0)
38        waitfortask.cancel()
39        await asyncio.wait({waitfortask})
40
41        self.assertTrue(t.exited)
42
43    def test_asyncio_wait_for_cancelled(self):
44        asyncio.run(self.atest_asyncio_wait_for_cancelled())
45
46    async def atest_asyncio_wait_for_timeout(self):
47        t  = SlowTask()
48
49        try:
50            await asyncio.wait_for(t.run(), t.TASK_TIMEOUT / 2)
51        except asyncio.TimeoutError:
52            pass
53
54        self.assertTrue(t.exited)
55
56    def test_asyncio_wait_for_timeout(self):
57        asyncio.run(self.atest_asyncio_wait_for_timeout())
58
59
60if __name__ == '__main__':
61    unittest.main()
62