| 631 | timeout = get_async_test_timeout() |
| 632 | |
| 633 | def wrap(f: Callable[..., Union[Generator, "Coroutine"]]) -> Callable[..., None]: |
| 634 | # Stack up several decorators to allow us to access the generator |
| 635 | # object itself. In the innermost wrapper, we capture the generator |
| 636 | # and save it in an attribute of self. Next, we run the wrapped |
| 637 | # function through @gen.coroutine. Finally, the coroutine is |
| 638 | # wrapped again to make it synchronous with run_sync. |
| 639 | # |
| 640 | # This is a good case study arguing for either some sort of |
| 641 | # extensibility in the gen decorators or cancellation support. |
| 642 | @functools.wraps(f) |
| 643 | def pre_coroutine(self, *args, **kwargs): |
| 644 | # type: (AsyncTestCase, *Any, **Any) -> Union[Generator, Coroutine] |
| 645 | # Type comments used to avoid pypy3 bug. |
| 646 | result = f(self, *args, **kwargs) |
| 647 | if isinstance(result, Generator) or inspect.iscoroutine(result): |
| 648 | self._test_generator = result |
| 649 | else: |
| 650 | self._test_generator = None |
| 651 | return result |
| 652 | |
| 653 | if inspect.iscoroutinefunction(f): |
| 654 | coro = pre_coroutine |
| 655 | else: |
| 656 | coro = gen.coroutine(pre_coroutine) # type: ignore[assignment] |
| 657 | |
| 658 | @functools.wraps(coro) |
| 659 | def post_coroutine(self, *args, **kwargs): |
| 660 | # type: (AsyncTestCase, *Any, **Any) -> None |
| 661 | try: |
| 662 | return self.io_loop.run_sync( |
| 663 | functools.partial(coro, self, *args, **kwargs), timeout=timeout |
| 664 | ) |
| 665 | except TimeoutError as e: |
| 666 | # run_sync raises an error with an unhelpful traceback. |
| 667 | # If the underlying generator is still running, we can throw the |
| 668 | # exception back into it so the stack trace is replaced by the |
| 669 | # point where the test is stopped. The only reason the generator |
| 670 | # would not be running would be if it were cancelled, which means |
| 671 | # a native coroutine, so we can rely on the cr_running attribute. |
| 672 | if self._test_generator is not None and getattr( |
| 673 | self._test_generator, "cr_running", True |
| 674 | ): |
| 675 | self._test_generator.throw(type(e), e) |
| 676 | # In case the test contains an overly broad except |
| 677 | # clause, we may get back here. |
| 678 | # Coroutine was stopped or didn't raise a useful stack trace, |
| 679 | # so re-raise the original exception which is better than nothing. |
| 680 | raise |
| 681 | |
| 682 | return post_coroutine |
| 683 | |
| 684 | if func is not None: |
| 685 | # Used like: |