diff options
| -rw-r--r-- | Lib/test/test_sys_settrace.py | 32 |
1 files changed, 28 insertions, 4 deletions
diff --git a/Lib/test/test_sys_settrace.py b/Lib/test/test_sys_settrace.py index 6fb9099..eef43e8 100644 --- a/Lib/test/test_sys_settrace.py +++ b/Lib/test/test_sys_settrace.py @@ -41,6 +41,20 @@ async def asynciter(iterable): for x in iterable: yield x +def clean_asynciter(test): + @wraps(test) + async def wrapper(*args, **kwargs): + cleanups = [] + def wrapped_asynciter(iterable): + it = asynciter(iterable) + cleanups.append(it.aclose) + return it + try: + return await test(*args, **kwargs, asynciter=wrapped_asynciter) + finally: + while cleanups: + await cleanups.pop()() + return wrapper # A very basic example. If this fails, we're in deep trouble. def basic(): @@ -1914,7 +1928,11 @@ class JumpTestCase(unittest.TestCase): def run_test(self, func, jumpFrom, jumpTo, expected, error=None, event='line', decorated=False, warning=None): - tracer = JumpTracer(func, jumpFrom, jumpTo, event, decorated) + wrapped = func + while hasattr(wrapped, '__wrapped__'): + wrapped = wrapped.__wrapped__ + + tracer = JumpTracer(wrapped, jumpFrom, jumpTo, event, decorated) sys.settrace(tracer.trace) output = [] @@ -1930,7 +1948,11 @@ class JumpTestCase(unittest.TestCase): def run_async_test(self, func, jumpFrom, jumpTo, expected, error=None, event='line', decorated=False, warning=None): - tracer = JumpTracer(func, jumpFrom, jumpTo, event, decorated) + wrapped = func + while hasattr(wrapped, '__wrapped__'): + wrapped = wrapped.__wrapped__ + + tracer = JumpTracer(wrapped, jumpFrom, jumpTo, event, decorated) sys.settrace(tracer.trace) output = [] @@ -2001,7 +2023,8 @@ class JumpTestCase(unittest.TestCase): output.append(7) @async_jump_test(4, 5, [3, 5]) - async def test_jump_out_of_async_for_block_forwards(output): + @clean_asynciter + async def test_jump_out_of_async_for_block_forwards(output, asynciter): for i in [1]: async for i in asynciter([1, 2]): output.append(3) @@ -2009,7 +2032,8 @@ class JumpTestCase(unittest.TestCase): output.append(5) @async_jump_test(5, 2, [2, 4, 2, 4, 5, 6]) - async def test_jump_out_of_async_for_block_backwards(output): + @clean_asynciter + async def test_jump_out_of_async_for_block_backwards(output, asynciter): for i in [1]: output.append(2) async for i in asynciter([1]): |