diff --git a/test/asynchronous/test_discovery_and_monitoring.py b/test/asynchronous/test_discovery_and_monitoring.py index 84382cf86f..1e2877498b 100644 --- a/test/asynchronous/test_discovery_and_monitoring.py +++ b/test/asynchronous/test_discovery_and_monitoring.py @@ -394,7 +394,7 @@ async def test_pool_unpause(self): @async_client_context.require_failCommand_appName @async_client_context.require_test_commands @async_client_context.require_async - @flaky(reason="PYTHON-5428") + @flaky(reason="PyPy is slow") async def test_connection_close_does_not_block_other_operations(self): listener = CMAPHeartbeatListener() client = await self.async_single_client( diff --git a/test/asynchronous/utils.py b/test/asynchronous/utils.py index 2b59d151eb..1407e38d91 100644 --- a/test/asynchronous/utils.py +++ b/test/asynchronous/utils.py @@ -24,6 +24,7 @@ import threading # Used in the synchronized version of this file import time import traceback +import unittest from functools import wraps from inspect import iscoroutinefunction @@ -150,6 +151,19 @@ async def async_joinall(tasks): await asyncio.wait([t.task for t in tasks if t is not None], timeout=300) +async def _run_attempt_cleanups(method, depth): + """Run the cleanups a failed flaky attempt registered, most recent first.""" + while len(method._cleanups) > depth: + func, args, kwargs = method._cleanups.pop() + try: + if iscoroutinefunction(func): + await func(*args, **kwargs) + else: + func(*args, **kwargs) + except Exception: + traceback.print_exc() + + def flaky( *, reason=None, @@ -162,6 +176,9 @@ def flaky( ): """Decorate a test as flaky. + Before each retry, any cleanups registered on the test case during the + failed attempt are run, then ``reset_func`` is called. + :param reason: the reason why the test is flaky :param max_runs: the maximum number of runs before raising an error :param min_passes: the minimum number of passing runs @@ -185,8 +202,13 @@ def flaky( def decorator(target_func): @wraps(target_func) async def wrapper(*args, **kwargs): + # flaky decorates either an unbound test method (prose test) or a bound method (unified test). + method = getattr(target_func, "__self__", None) + if method is None and args and isinstance(args[0], unittest.TestCase): + method = args[0] passes = 0 for i in range(max_runs): + depth = len(method._cleanups) if method is not None else 0 try: result = await target_func(*args, **kwargs) passes += 1 @@ -200,6 +222,8 @@ async def wrapper(*args, **kwargs): f"{traceback.format_exc()}", file=sys.stderr, ) + if method is not None: + await _run_attempt_cleanups(method, depth) await asyncio.sleep(delay) if reset_func: await reset_func() diff --git a/test/test_discovery_and_monitoring.py b/test/test_discovery_and_monitoring.py index 7b36cc318a..72098deea2 100644 --- a/test/test_discovery_and_monitoring.py +++ b/test/test_discovery_and_monitoring.py @@ -394,7 +394,7 @@ def test_pool_unpause(self): @client_context.require_failCommand_appName @client_context.require_test_commands @client_context.require_async - @flaky(reason="PYTHON-5428") + @flaky(reason="PyPy is slow") def test_connection_close_does_not_block_other_operations(self): listener = CMAPHeartbeatListener() client = self.single_client( diff --git a/test/utils.py b/test/utils.py index 1526e2a828..55af404c32 100644 --- a/test/utils.py +++ b/test/utils.py @@ -24,6 +24,7 @@ import threading # Used in the synchronized version of this file import time import traceback +import unittest from functools import wraps from inspect import iscoroutinefunction @@ -150,6 +151,19 @@ def joinall(tasks): asyncio.wait([t.task for t in tasks if t is not None], timeout=300) +def _run_attempt_cleanups(method, depth): + """Run the cleanups a failed flaky attempt registered, most recent first.""" + while len(method._cleanups) > depth: + func, args, kwargs = method._cleanups.pop() + try: + if iscoroutinefunction(func): + func(*args, **kwargs) + else: + func(*args, **kwargs) + except Exception: + traceback.print_exc() + + def flaky( *, reason=None, @@ -162,6 +176,9 @@ def flaky( ): """Decorate a test as flaky. + Before each retry, any cleanups registered on the test case during the + failed attempt are run, then ``reset_func`` is called. + :param reason: the reason why the test is flaky :param max_runs: the maximum number of runs before raising an error :param min_passes: the minimum number of passing runs @@ -185,8 +202,13 @@ def flaky( def decorator(target_func): @wraps(target_func) def wrapper(*args, **kwargs): + # flaky decorates either an unbound test method (prose test) or a bound method (unified test). + method = getattr(target_func, "__self__", None) + if method is None and args and isinstance(args[0], unittest.TestCase): + method = args[0] passes = 0 for i in range(max_runs): + depth = len(method._cleanups) if method is not None else 0 try: result = target_func(*args, **kwargs) passes += 1 @@ -200,6 +222,8 @@ def wrapper(*args, **kwargs): f"{traceback.format_exc()}", file=sys.stderr, ) + if method is not None: + _run_attempt_cleanups(method, depth) time.sleep(delay) if reset_func: reset_func()