blob: 41b1b4f4b48547614600cf6a2f2785c2ad1cfc08 [file] [log] [blame]
Yury Selivanoveb636452016-09-08 22:01:51 -07001import asyncio
2import inspect
3import sys
4import types
5import unittest
6
7from unittest import mock
8
9
10class AwaitException(Exception):
11 pass
12
13
14@types.coroutine
15def awaitable(*, throw=False):
16 if throw:
17 yield ('throw',)
18 else:
19 yield ('result',)
20
21
22def run_until_complete(coro):
23 exc = False
24 while True:
25 try:
26 if exc:
27 exc = False
28 fut = coro.throw(AwaitException)
29 else:
30 fut = coro.send(None)
31 except StopIteration as ex:
32 return ex.args[0]
33
34 if fut == ('throw',):
35 exc = True
36
37
38def to_list(gen):
39 async def iterate():
40 res = []
41 async for i in gen:
42 res.append(i)
43 return res
44
45 return run_until_complete(iterate())
46
47
48class AsyncGenSyntaxTest(unittest.TestCase):
49
50 def test_async_gen_syntax_01(self):
51 code = '''async def foo():
52 await abc
53 yield from 123
54 '''
55
56 with self.assertRaisesRegex(SyntaxError, 'yield from.*inside async'):
57 exec(code, {}, {})
58
59 def test_async_gen_syntax_02(self):
60 code = '''async def foo():
61 yield from 123
62 '''
63
64 with self.assertRaisesRegex(SyntaxError, 'yield from.*inside async'):
65 exec(code, {}, {})
66
67 def test_async_gen_syntax_03(self):
68 code = '''async def foo():
69 await abc
70 yield
71 return 123
72 '''
73
74 with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'):
75 exec(code, {}, {})
76
77 def test_async_gen_syntax_04(self):
78 code = '''async def foo():
79 yield
80 return 123
81 '''
82
83 with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'):
84 exec(code, {}, {})
85
86 def test_async_gen_syntax_05(self):
87 code = '''async def foo():
88 if 0:
89 yield
90 return 12
91 '''
92
93 with self.assertRaisesRegex(SyntaxError, 'return.*value.*async gen'):
94 exec(code, {}, {})
95
96
97class AsyncGenTest(unittest.TestCase):
98
99 def compare_generators(self, sync_gen, async_gen):
100 def sync_iterate(g):
101 res = []
102 while True:
103 try:
104 res.append(g.__next__())
105 except StopIteration:
106 res.append('STOP')
107 break
108 except Exception as ex:
109 res.append(str(type(ex)))
110 return res
111
112 def async_iterate(g):
113 res = []
114 while True:
115 try:
116 g.__anext__().__next__()
117 except StopAsyncIteration:
118 res.append('STOP')
119 break
120 except StopIteration as ex:
121 if ex.args:
122 res.append(ex.args[0])
123 else:
124 res.append('EMPTY StopIteration')
125 break
126 except Exception as ex:
127 res.append(str(type(ex)))
128 return res
129
130 sync_gen_result = sync_iterate(sync_gen)
131 async_gen_result = async_iterate(async_gen)
132 self.assertEqual(sync_gen_result, async_gen_result)
133 return async_gen_result
134
135 def test_async_gen_iteration_01(self):
136 async def gen():
137 await awaitable()
138 a = yield 123
139 self.assertIs(a, None)
140 await awaitable()
141 yield 456
142 await awaitable()
143 yield 789
144
145 self.assertEqual(to_list(gen()), [123, 456, 789])
146
147 def test_async_gen_iteration_02(self):
148 async def gen():
149 await awaitable()
150 yield 123
151 await awaitable()
152
153 g = gen()
154 ai = g.__aiter__()
155 self.assertEqual(ai.__anext__().__next__(), ('result',))
156
157 try:
158 ai.__anext__().__next__()
159 except StopIteration as ex:
160 self.assertEqual(ex.args[0], 123)
161 else:
162 self.fail('StopIteration was not raised')
163
164 self.assertEqual(ai.__anext__().__next__(), ('result',))
165
166 try:
167 ai.__anext__().__next__()
168 except StopAsyncIteration as ex:
169 self.assertFalse(ex.args)
170 else:
171 self.fail('StopAsyncIteration was not raised')
172
173 def test_async_gen_exception_03(self):
174 async def gen():
175 await awaitable()
176 yield 123
177 await awaitable(throw=True)
178 yield 456
179
180 with self.assertRaises(AwaitException):
181 to_list(gen())
182
183 def test_async_gen_exception_04(self):
184 async def gen():
185 await awaitable()
186 yield 123
187 1 / 0
188
189 g = gen()
190 ai = g.__aiter__()
191 self.assertEqual(ai.__anext__().__next__(), ('result',))
192
193 try:
194 ai.__anext__().__next__()
195 except StopIteration as ex:
196 self.assertEqual(ex.args[0], 123)
197 else:
198 self.fail('StopIteration was not raised')
199
200 with self.assertRaises(ZeroDivisionError):
201 ai.__anext__().__next__()
202
203 def test_async_gen_exception_05(self):
204 async def gen():
205 yield 123
206 raise StopAsyncIteration
207
208 with self.assertRaisesRegex(RuntimeError,
209 'async generator.*StopAsyncIteration'):
210 to_list(gen())
211
212 def test_async_gen_exception_06(self):
213 async def gen():
214 yield 123
215 raise StopIteration
216
217 with self.assertRaisesRegex(RuntimeError,
218 'async generator.*StopIteration'):
219 to_list(gen())
220
221 def test_async_gen_exception_07(self):
222 def sync_gen():
223 try:
224 yield 1
225 1 / 0
226 finally:
227 yield 2
228 yield 3
229
230 yield 100
231
232 async def async_gen():
233 try:
234 yield 1
235 1 / 0
236 finally:
237 yield 2
238 yield 3
239
240 yield 100
241
242 self.compare_generators(sync_gen(), async_gen())
243
244 def test_async_gen_exception_08(self):
245 def sync_gen():
246 try:
247 yield 1
248 finally:
249 yield 2
250 1 / 0
251 yield 3
252
253 yield 100
254
255 async def async_gen():
256 try:
257 yield 1
258 await awaitable()
259 finally:
260 await awaitable()
261 yield 2
262 1 / 0
263 yield 3
264
265 yield 100
266
267 self.compare_generators(sync_gen(), async_gen())
268
269 def test_async_gen_exception_09(self):
270 def sync_gen():
271 try:
272 yield 1
273 1 / 0
274 finally:
275 yield 2
276 yield 3
277
278 yield 100
279
280 async def async_gen():
281 try:
282 await awaitable()
283 yield 1
284 1 / 0
285 finally:
286 yield 2
287 await awaitable()
288 yield 3
289
290 yield 100
291
292 self.compare_generators(sync_gen(), async_gen())
293
294 def test_async_gen_exception_10(self):
295 async def gen():
296 yield 123
297 with self.assertRaisesRegex(TypeError,
298 "non-None value .* async generator"):
299 gen().__anext__().send(100)
300
301 def test_async_gen_api_01(self):
302 async def gen():
303 yield 123
304
305 g = gen()
306
307 self.assertEqual(g.__name__, 'gen')
308 g.__name__ = '123'
309 self.assertEqual(g.__name__, '123')
310
311 self.assertIn('.gen', g.__qualname__)
312 g.__qualname__ = '123'
313 self.assertEqual(g.__qualname__, '123')
314
315 self.assertIsNone(g.ag_await)
316 self.assertIsInstance(g.ag_frame, types.FrameType)
317 self.assertFalse(g.ag_running)
318 self.assertIsInstance(g.ag_code, types.CodeType)
319
320 self.assertTrue(inspect.isawaitable(g.aclose()))
321
322
323class AsyncGenAsyncioTest(unittest.TestCase):
324
325 def setUp(self):
326 self.loop = asyncio.new_event_loop()
327 asyncio.set_event_loop(None)
328
329 def tearDown(self):
330 self.loop.close()
331 self.loop = None
332
333 async def to_list(self, gen):
334 res = []
335 async for i in gen:
336 res.append(i)
337 return res
338
339 def test_async_gen_asyncio_01(self):
340 async def gen():
341 yield 1
342 await asyncio.sleep(0.01, loop=self.loop)
343 yield 2
344 await asyncio.sleep(0.01, loop=self.loop)
345 return
346 yield 3
347
348 res = self.loop.run_until_complete(self.to_list(gen()))
349 self.assertEqual(res, [1, 2])
350
351 def test_async_gen_asyncio_02(self):
352 async def gen():
353 yield 1
354 await asyncio.sleep(0.01, loop=self.loop)
355 yield 2
356 1 / 0
357 yield 3
358
359 with self.assertRaises(ZeroDivisionError):
360 self.loop.run_until_complete(self.to_list(gen()))
361
362 def test_async_gen_asyncio_03(self):
363 loop = self.loop
364
365 class Gen:
366 async def __aiter__(self):
367 yield 1
368 await asyncio.sleep(0.01, loop=loop)
369 yield 2
370
371 res = loop.run_until_complete(self.to_list(Gen()))
372 self.assertEqual(res, [1, 2])
373
374 def test_async_gen_asyncio_anext_04(self):
375 async def foo():
376 yield 1
377 await asyncio.sleep(0.01, loop=self.loop)
378 try:
379 yield 2
380 yield 3
381 except ZeroDivisionError:
382 yield 1000
383 await asyncio.sleep(0.01, loop=self.loop)
384 yield 4
385
386 async def run1():
387 it = foo().__aiter__()
388
389 self.assertEqual(await it.__anext__(), 1)
390 self.assertEqual(await it.__anext__(), 2)
391 self.assertEqual(await it.__anext__(), 3)
392 self.assertEqual(await it.__anext__(), 4)
393 with self.assertRaises(StopAsyncIteration):
394 await it.__anext__()
395 with self.assertRaises(StopAsyncIteration):
396 await it.__anext__()
397
398 async def run2():
399 it = foo().__aiter__()
400
401 self.assertEqual(await it.__anext__(), 1)
402 self.assertEqual(await it.__anext__(), 2)
403 try:
404 it.__anext__().throw(ZeroDivisionError)
405 except StopIteration as ex:
406 self.assertEqual(ex.args[0], 1000)
407 else:
408 self.fail('StopIteration was not raised')
409 self.assertEqual(await it.__anext__(), 4)
410 with self.assertRaises(StopAsyncIteration):
411 await it.__anext__()
412
413 self.loop.run_until_complete(run1())
414 self.loop.run_until_complete(run2())
415
416 def test_async_gen_asyncio_anext_05(self):
417 async def foo():
418 v = yield 1
419 v = yield v
420 yield v * 100
421
422 async def run():
423 it = foo().__aiter__()
424
425 try:
426 it.__anext__().send(None)
427 except StopIteration as ex:
428 self.assertEqual(ex.args[0], 1)
429 else:
430 self.fail('StopIteration was not raised')
431
432 try:
433 it.__anext__().send(10)
434 except StopIteration as ex:
435 self.assertEqual(ex.args[0], 10)
436 else:
437 self.fail('StopIteration was not raised')
438
439 try:
440 it.__anext__().send(12)
441 except StopIteration as ex:
442 self.assertEqual(ex.args[0], 1200)
443 else:
444 self.fail('StopIteration was not raised')
445
446 with self.assertRaises(StopAsyncIteration):
447 await it.__anext__()
448
449 self.loop.run_until_complete(run())
450
451 def test_async_gen_asyncio_aclose_06(self):
452 async def foo():
453 try:
454 yield 1
455 1 / 0
456 finally:
457 await asyncio.sleep(0.01, loop=self.loop)
458 yield 12
459
460 async def run():
461 gen = foo()
462 it = gen.__aiter__()
463 await it.__anext__()
464 await gen.aclose()
465
466 with self.assertRaisesRegex(
467 RuntimeError,
468 "async generator ignored GeneratorExit"):
469 self.loop.run_until_complete(run())
470
471 def test_async_gen_asyncio_aclose_07(self):
472 DONE = 0
473
474 async def foo():
475 nonlocal DONE
476 try:
477 yield 1
478 1 / 0
479 finally:
480 await asyncio.sleep(0.01, loop=self.loop)
481 await asyncio.sleep(0.01, loop=self.loop)
482 DONE += 1
483 DONE += 1000
484
485 async def run():
486 gen = foo()
487 it = gen.__aiter__()
488 await it.__anext__()
489 await gen.aclose()
490
491 self.loop.run_until_complete(run())
492 self.assertEqual(DONE, 1)
493
494 def test_async_gen_asyncio_aclose_08(self):
495 DONE = 0
496
497 fut = asyncio.Future(loop=self.loop)
498
499 async def foo():
500 nonlocal DONE
501 try:
502 yield 1
503 await fut
504 DONE += 1000
505 yield 2
506 finally:
507 await asyncio.sleep(0.01, loop=self.loop)
508 await asyncio.sleep(0.01, loop=self.loop)
509 DONE += 1
510 DONE += 1000
511
512 async def run():
513 gen = foo()
514 it = gen.__aiter__()
515 self.assertEqual(await it.__anext__(), 1)
516 t = self.loop.create_task(it.__anext__())
517 await asyncio.sleep(0.01, loop=self.loop)
518 await gen.aclose()
519 return t
520
521 t = self.loop.run_until_complete(run())
522 self.assertEqual(DONE, 1)
523
524 # Silence ResourceWarnings
525 fut.cancel()
526 t.cancel()
527 self.loop.run_until_complete(asyncio.sleep(0.01, loop=self.loop))
528
529 def test_async_gen_asyncio_gc_aclose_09(self):
530 DONE = 0
531
532 async def gen():
533 nonlocal DONE
534 try:
535 while True:
536 yield 1
537 finally:
538 await asyncio.sleep(0.01, loop=self.loop)
539 await asyncio.sleep(0.01, loop=self.loop)
540 DONE = 1
541
542 async def run():
543 g = gen()
544 await g.__anext__()
545 await g.__anext__()
546 del g
547
548 await asyncio.sleep(0.1, loop=self.loop)
549
550 self.loop.run_until_complete(run())
551 self.assertEqual(DONE, 1)
552
553 def test_async_gen_asyncio_asend_01(self):
554 DONE = 0
555
556 # Sanity check:
557 def sgen():
558 v = yield 1
559 yield v * 2
560 sg = sgen()
561 v = sg.send(None)
562 self.assertEqual(v, 1)
563 v = sg.send(100)
564 self.assertEqual(v, 200)
565
566 async def gen():
567 nonlocal DONE
568 try:
569 await asyncio.sleep(0.01, loop=self.loop)
570 v = yield 1
571 await asyncio.sleep(0.01, loop=self.loop)
572 yield v * 2
573 await asyncio.sleep(0.01, loop=self.loop)
574 return
575 finally:
576 await asyncio.sleep(0.01, loop=self.loop)
577 await asyncio.sleep(0.01, loop=self.loop)
578 DONE = 1
579
580 async def run():
581 g = gen()
582
583 v = await g.asend(None)
584 self.assertEqual(v, 1)
585
586 v = await g.asend(100)
587 self.assertEqual(v, 200)
588
589 with self.assertRaises(StopAsyncIteration):
590 await g.asend(None)
591
592 self.loop.run_until_complete(run())
593 self.assertEqual(DONE, 1)
594
595 def test_async_gen_asyncio_asend_02(self):
596 DONE = 0
597
598 async def sleep_n_crash(delay):
599 await asyncio.sleep(delay, loop=self.loop)
600 1 / 0
601
602 async def gen():
603 nonlocal DONE
604 try:
605 await asyncio.sleep(0.01, loop=self.loop)
606 v = yield 1
607 await sleep_n_crash(0.01)
608 DONE += 1000
609 yield v * 2
610 finally:
611 await asyncio.sleep(0.01, loop=self.loop)
612 await asyncio.sleep(0.01, loop=self.loop)
613 DONE = 1
614
615 async def run():
616 g = gen()
617
618 v = await g.asend(None)
619 self.assertEqual(v, 1)
620
621 await g.asend(100)
622
623 with self.assertRaises(ZeroDivisionError):
624 self.loop.run_until_complete(run())
625 self.assertEqual(DONE, 1)
626
627 def test_async_gen_asyncio_asend_03(self):
628 DONE = 0
629
630 async def sleep_n_crash(delay):
631 fut = asyncio.ensure_future(asyncio.sleep(delay, loop=self.loop),
632 loop=self.loop)
633 self.loop.call_later(delay / 2, lambda: fut.cancel())
634 return await fut
635
636 async def gen():
637 nonlocal DONE
638 try:
639 await asyncio.sleep(0.01, loop=self.loop)
640 v = yield 1
641 await sleep_n_crash(0.01)
642 DONE += 1000
643 yield v * 2
644 finally:
645 await asyncio.sleep(0.01, loop=self.loop)
646 await asyncio.sleep(0.01, loop=self.loop)
647 DONE = 1
648
649 async def run():
650 g = gen()
651
652 v = await g.asend(None)
653 self.assertEqual(v, 1)
654
655 await g.asend(100)
656
657 with self.assertRaises(asyncio.CancelledError):
658 self.loop.run_until_complete(run())
659 self.assertEqual(DONE, 1)
660
661 def test_async_gen_asyncio_athrow_01(self):
662 DONE = 0
663
664 class FooEr(Exception):
665 pass
666
667 # Sanity check:
668 def sgen():
669 try:
670 v = yield 1
671 except FooEr:
672 v = 1000
673 yield v * 2
674 sg = sgen()
675 v = sg.send(None)
676 self.assertEqual(v, 1)
677 v = sg.throw(FooEr)
678 self.assertEqual(v, 2000)
679 with self.assertRaises(StopIteration):
680 sg.send(None)
681
682 async def gen():
683 nonlocal DONE
684 try:
685 await asyncio.sleep(0.01, loop=self.loop)
686 try:
687 v = yield 1
688 except FooEr:
689 v = 1000
690 await asyncio.sleep(0.01, loop=self.loop)
691 yield v * 2
692 await asyncio.sleep(0.01, loop=self.loop)
693 # return
694 finally:
695 await asyncio.sleep(0.01, loop=self.loop)
696 await asyncio.sleep(0.01, loop=self.loop)
697 DONE = 1
698
699 async def run():
700 g = gen()
701
702 v = await g.asend(None)
703 self.assertEqual(v, 1)
704
705 v = await g.athrow(FooEr)
706 self.assertEqual(v, 2000)
707
708 with self.assertRaises(StopAsyncIteration):
709 await g.asend(None)
710
711 self.loop.run_until_complete(run())
712 self.assertEqual(DONE, 1)
713
714 def test_async_gen_asyncio_athrow_02(self):
715 DONE = 0
716
717 class FooEr(Exception):
718 pass
719
720 async def sleep_n_crash(delay):
721 fut = asyncio.ensure_future(asyncio.sleep(delay, loop=self.loop),
722 loop=self.loop)
723 self.loop.call_later(delay / 2, lambda: fut.cancel())
724 return await fut
725
726 async def gen():
727 nonlocal DONE
728 try:
729 await asyncio.sleep(0.01, loop=self.loop)
730 try:
731 v = yield 1
732 except FooEr:
733 await sleep_n_crash(0.01)
734 yield v * 2
735 await asyncio.sleep(0.01, loop=self.loop)
736 # return
737 finally:
738 await asyncio.sleep(0.01, loop=self.loop)
739 await asyncio.sleep(0.01, loop=self.loop)
740 DONE = 1
741
742 async def run():
743 g = gen()
744
745 v = await g.asend(None)
746 self.assertEqual(v, 1)
747
748 try:
749 await g.athrow(FooEr)
750 except asyncio.CancelledError:
751 self.assertEqual(DONE, 1)
752 raise
753 else:
754 self.fail('CancelledError was not raised')
755
756 with self.assertRaises(asyncio.CancelledError):
757 self.loop.run_until_complete(run())
758 self.assertEqual(DONE, 1)
759
760 def test_async_gen_asyncio_shutdown_01(self):
761 finalized = 0
762
763 async def waiter(timeout):
764 nonlocal finalized
765 try:
766 await asyncio.sleep(timeout, loop=self.loop)
767 yield 1
768 finally:
769 await asyncio.sleep(0, loop=self.loop)
770 finalized += 1
771
772 async def wait():
773 async for _ in waiter(1):
774 pass
775
776 t1 = self.loop.create_task(wait())
777 t2 = self.loop.create_task(wait())
778
779 self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop))
780
781 self.loop.run_until_complete(self.loop.shutdown_asyncgens())
782 self.assertEqual(finalized, 2)
783
784 # Silence warnings
785 t1.cancel()
786 t2.cancel()
787 self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop))
788
789 def test_async_gen_asyncio_shutdown_02(self):
790 logged = 0
791
792 def logger(loop, context):
793 nonlocal logged
794 self.assertIn('asyncgen', context)
795 expected = 'an error occurred during closing of asynchronous'
796 if expected in context['message']:
797 logged += 1
798
799 async def waiter(timeout):
800 try:
801 await asyncio.sleep(timeout, loop=self.loop)
802 yield 1
803 finally:
804 1 / 0
805
806 async def wait():
807 async for _ in waiter(1):
808 pass
809
810 t = self.loop.create_task(wait())
811 self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop))
812
813 self.loop.set_exception_handler(logger)
814 self.loop.run_until_complete(self.loop.shutdown_asyncgens())
815
816 self.assertEqual(logged, 1)
817
818 # Silence warnings
819 t.cancel()
820 self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop))
821
822if __name__ == "__main__":
823 unittest.main()