clean up flaky test handling

Author: blink1073

test/asynchronous/unified_format.py

@@ -35,7 +35,7 @@
     client_knobs,
     unittest,
 )
-from test.asynchronous.utils import async_get_pool
+from test.asynchronous.utils import async_get_pool, flaky
 from test.asynchronous.utils_spec_runner import SpecRunnerTask
 from test.unified_format_shared import (
     KMS_TLS_OPTS,
@@ -529,20 +529,10 @@ def maybe_skip_test(self, spec):
             self.skipTest("Implement PYTHON-1894")
         if "timeoutMS applied to entire download" in spec["description"]:
             self.skipTest("PyMongo's open_download_stream does not cap the stream's lifetime")
-        if (
-            "Error returned from connection pool clear with interruptInUseConnections=true is retryable"
-            in spec["description"]
-            and not _IS_SYNC
-        ):
-            self.skipTest("PYTHON-5170 tests are flakey")
-        if "Driver extends timeout while streaming" in spec["description"] and not _IS_SYNC:
-            self.skipTest("PYTHON-5174 tests are flakey")
 
         class_name = self.__class__.__name__.lower()
         description = spec["description"].lower()
         if "csot" in class_name:
-            if "gridfs" in class_name and sys.platform == "win32":
-                self.skipTest("PYTHON-3522 CSOT GridFS tests are flaky on Windows")
             if async_client_context.storage_engine == "mmapv1":
                 self.skipTest(
                     "MMAPv1 does not support retryable writes which is required for CSOT tests"
@@ -1378,30 +1368,21 @@ async def run_scenario(self, spec, uri=None):
         # operations during test set up and tear down.
         await self.kill_all_sessions()
 
-        if "csot" in self.id().lower():
-            # Retry CSOT tests up to 2 times to deal with flakey tests.
-            attempts = 3
-            for i in range(attempts):
-                try:
-                    return await self._run_scenario(spec, uri)
-                except (AssertionError, OperationFailure) as exc:
-                    if isinstance(exc, OperationFailure) and (
-                        _IS_SYNC or "failpoint" not in exc._message
-                    ):
-                        raise
-                    if i < attempts - 1:
-                        print(
-                            f"Retrying after attempt {i+1} of {self.id()} failed with:\n"
-                            f"{traceback.format_exc()}",
-                            file=sys.stderr,
-                        )
-                        await self.asyncSetUp()
-                        continue
-                    raise
-            return None
-        else:
-            await self._run_scenario(spec, uri)
-            return None
+        # Handle flaky tests.
+        func = functools.partial(self._run_scenario, spec, uri)
+        flaky_tests = [
+            # PYTHON-5170
+            ".*test_discovery_and_monitoring.TestUnifiedInterruptInUsePoolClear.*",
+            # PYTHON-5174
+            ".*Driver_extends_timeout_while_streaming",
+            ".*csot.*",
+        ]
+        for flaky_test in flaky_tests:
+            if re.match(flaky_test, self.id()) is not None:
+                decorator = flaky(reset_func=self.asyncSetUp, func_name=self.id())
+                func = decorator(func)
+                break
+        await func()
 
     async def _run_scenario(self, spec, uri=None):
         # maybe skip test manually

test/asynchronous/utils.py

@@ -21,6 +21,7 @@
 import sys
 import threading  # Used in the synchronized version of this file
 import time
+import traceback
 from asyncio import iscoroutinefunction
 from functools import wraps
 
@@ -156,7 +157,16 @@ async def async_joinall(tasks):
         await asyncio.wait([t.task for t in tasks if t is not None], timeout=300)
 
 
-def flaky(func=None, *, max_runs=2, min_passes=1, delay=1, affects_cpython_linux=False):
+def flaky(
+    func=None,
+    *,
+    max_runs=2,
+    min_passes=1,
+    delay=1,
+    affects_cpython_linux=False,
+    func_name=None,
+    reset_func=None,
+):
     is_cpython_linux = sys.platform == "linux" and sys.implementation.name == "cpython"
     if is_cpython_linux and not affects_cpython_linux:
         max_runs = 1
@@ -177,7 +187,13 @@ async def wrapper(*args, **kwargs):
                     failure = e
                     await asyncio.sleep(delay)
                 if failure is not None and i < max_runs - 1:
-                    print("Flaky failure:", failure)
+                    print(
+                        f"Retrying after attempt {i+1} of {func_name or target_func.__name__} failed with:\n"
+                        f"{traceback.format_exc()}",
+                        file=sys.stderr,
+                    )
+                    if reset_func:
+                        await reset_func()
             if failure:
                 raise failure
             raise RuntimeError(f"Only passed {passes} of {min_passes} times")

test/unified_format.py

@@ -47,7 +47,7 @@
     parse_collection_or_database_options,
     with_metaclass,
 )
-from test.utils import get_pool
+from test.utils import flaky, get_pool
 from test.utils_shared import (
     camel_to_snake,
     camel_to_snake_args,
@@ -528,20 +528,10 @@ def maybe_skip_test(self, spec):
             self.skipTest("Implement PYTHON-1894")
         if "timeoutMS applied to entire download" in spec["description"]:
             self.skipTest("PyMongo's open_download_stream does not cap the stream's lifetime")
-        if (
-            "Error returned from connection pool clear with interruptInUseConnections=true is retryable"
-            in spec["description"]
-            and not _IS_SYNC
-        ):
-            self.skipTest("PYTHON-5170 tests are flakey")
-        if "Driver extends timeout while streaming" in spec["description"] and not _IS_SYNC:
-            self.skipTest("PYTHON-5174 tests are flakey")
 
         class_name = self.__class__.__name__.lower()
         description = spec["description"].lower()
         if "csot" in class_name:
-            if "gridfs" in class_name and sys.platform == "win32":
-                self.skipTest("PYTHON-3522 CSOT GridFS tests are flaky on Windows")
             if client_context.storage_engine == "mmapv1":
                 self.skipTest(
                     "MMAPv1 does not support retryable writes which is required for CSOT tests"
@@ -1365,30 +1355,21 @@ def run_scenario(self, spec, uri=None):
         # operations during test set up and tear down.
         self.kill_all_sessions()
 
-        if "csot" in self.id().lower():
-            # Retry CSOT tests up to 2 times to deal with flakey tests.
-            attempts = 3
-            for i in range(attempts):
-                try:
-                    return self._run_scenario(spec, uri)
-                except (AssertionError, OperationFailure) as exc:
-                    if isinstance(exc, OperationFailure) and (
-                        _IS_SYNC or "failpoint" not in exc._message
-                    ):
-                        raise
-                    if i < attempts - 1:
-                        print(
-                            f"Retrying after attempt {i+1} of {self.id()} failed with:\n"
-                            f"{traceback.format_exc()}",
-                            file=sys.stderr,
-                        )
-                        self.setUp()
-                        continue
-                    raise
-            return None
-        else:
-            self._run_scenario(spec, uri)
-            return None
+        # Handle flaky tests.
+        func = functools.partial(self._run_scenario, spec, uri)
+        flaky_tests = [
+            # PYTHON-5170
+            ".*test_discovery_and_monitoring.TestUnifiedInterruptInUsePoolClear.*",
+            # PYTHON-5174
+            ".*Driver_extends_timeout_while_streaming",
+            ".*csot.*",
+        ]
+        for flaky_test in flaky_tests:
+            if re.match(flaky_test, self.id()) is not None:
+                decorator = flaky(reset_func=self.setUp, func_name=self.id())
+                func = decorator(func)
+                break
+        func()
 
     def _run_scenario(self, spec, uri=None):
         # maybe skip test manually

test/utils.py

@@ -21,6 +21,7 @@
 import sys
 import threading  # Used in the synchronized version of this file
 import time
+import traceback
 from asyncio import iscoroutinefunction
 from functools import wraps
 
@@ -154,7 +155,16 @@ def joinall(tasks):
         asyncio.wait([t.task for t in tasks if t is not None], timeout=300)
 
 
-def flaky(func=None, *, max_runs=2, min_passes=1, delay=1, affects_cpython_linux=False):
+def flaky(
+    func=None,
+    *,
+    max_runs=2,
+    min_passes=1,
+    delay=1,
+    affects_cpython_linux=False,
+    func_name=None,
+    reset_func=None,
+):
     is_cpython_linux = sys.platform == "linux" and sys.implementation.name == "cpython"
     if is_cpython_linux and not affects_cpython_linux:
         max_runs = 1
@@ -175,7 +185,13 @@ def wrapper(*args, **kwargs):
                     failure = e
                     time.sleep(delay)
                 if failure is not None and i < max_runs - 1:
-                    print("Flaky failure:", failure)
+                    print(
+                        f"Retrying after attempt {i+1} of {func_name or target_func.__name__} failed with:\n"
+                        f"{traceback.format_exc()}",
+                        file=sys.stderr,
+                    )
+                    if reset_func:
+                        reset_func()
             if failure:
                 raise failure
             raise RuntimeError(f"Only passed {passes} of {min_passes} times")