PYTHON-5428 - @flaky calls cleanup between attempts

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(

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()

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(

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()