PYTHON-5413 Handle flaky tests

Author: blink1073

.evergreen/scripts/setup_tests.py

@@ -162,10 +162,6 @@ def handle_test_env() -> None:
     write_env("PIP_PREFER_BINARY")  # Prefer binary dists by default.
     write_env("UV_FROZEN")  # Do not modify lock files.
 
-    # Skip CSOT tests on non-linux platforms.
-    if PLATFORM != "linux":
-        write_env("SKIP_CSOT_TESTS")
-
     # Set an environment variable for the test name and sub test name.
     write_env(f"TEST_{test_name.upper()}")
     write_env("TEST_NAME", test_name)

test/asynchronous/test_client_bulk_write.py

@@ -25,6 +25,7 @@
     async_client_context,
     unittest,
 )
+from test.asynchronous.utils import flaky
 from test.utils_shared import (
     OvertCommandListener,
 )
@@ -619,15 +620,14 @@ async def test_15_unacknowledged_write_across_batches(self):
 # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#11-multi-batch-bulkwrites
 class TestClientBulkWriteCSOT(AsyncIntegrationTest):
     async def asyncSetUp(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         await super().asyncSetUp()
         self.max_write_batch_size = await async_client_context.max_write_batch_size
         self.max_bson_object_size = await async_client_context.max_bson_size
         self.max_message_size_bytes = await async_client_context.max_message_size_bytes
 
     @async_client_context.require_version_min(8, 0, 0, -24)
     @async_client_context.require_failCommand_fail_point
+    @flaky
     async def test_timeout_in_multi_batch_bulk_write(self):
         _OVERHEAD = 500
 

test/asynchronous/test_csot.py

@@ -23,6 +23,7 @@
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
 from test.asynchronous.unified_format import generate_test_classes
+from test.asynchronous.utils import flaky
 
 import pymongo
 from pymongo import _csot
@@ -43,9 +44,8 @@
 class TestCSOT(AsyncIntegrationTest):
     RUN_ON_LOAD_BALANCER = True
 
+    @flaky
     async def test_timeout_nested(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         coll = self.db.coll
         self.assertEqual(_csot.get_timeout(), None)
         self.assertEqual(_csot.get_deadline(), float("inf"))
@@ -82,9 +82,8 @@ async def test_timeout_nested(self):
         self.assertEqual(_csot.get_rtt(), 0.0)
 
     @async_client_context.require_change_streams
+    @flaky
     async def test_change_stream_can_resume_after_timeouts(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         coll = self.db.test
         await coll.insert_one({})
         async with await coll.watch() as stream:

test/asynchronous/test_cursor.py

@@ -31,6 +31,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
+from test.asynchronous.utils import flaky
 from test.utils_shared import (
     AllowListEventListener,
     EventListener,
@@ -1415,9 +1416,8 @@ async def test_to_list_length(self):
         docs = await c.to_list(3)
         self.assertEqual(len(docs), 2)
 
+    @flaky
     async def test_to_list_csot_applied(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         client = await self.async_single_client(timeoutMS=500, w=1)
         coll = client.pymongo.test
         # Initialize the client with a larger timeout to help make test less flakey
@@ -1458,9 +1458,8 @@ async def test_command_cursor_to_list_length(self):
         self.assertEqual(len(await result.to_list(1)), 1)
 
     @async_client_context.require_failCommand_blockConnection
+    @flaky
     async def test_command_cursor_to_list_csot_applied(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         client = await self.async_single_client(timeoutMS=500, w=1)
         coll = client.pymongo.test
         # Initialize the client with a larger timeout to help make test less flakey

test/asynchronous/unified_format.py

@@ -40,7 +40,6 @@
 from test.unified_format_shared import (
     KMS_TLS_OPTS,
     PLACEHOLDER_MAP,
-    SKIP_CSOT_TESTS,
     EventListenerUtil,
     MatchEvaluatorUtil,
     coerce_result,
@@ -1374,9 +1373,6 @@ async def verify_outcome(self, spec):
                 self.assertListEqual(sorted_expected_documents, actual_documents)
 
     async def run_scenario(self, spec, uri=None):
-        if "csot" in self.id().lower() and SKIP_CSOT_TESTS:
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
-
         # Kill all sessions before and after each test to prevent an open
         # transaction (from a test failure) from blocking collection/database
         # operations during test set up and tear down.

test/asynchronous/utils.py

@@ -18,9 +18,11 @@
 import asyncio
 import contextlib
 import random
+import sys
 import threading  # Used in the synchronized version of this file
 import time
 from asyncio import iscoroutinefunction
+from functools import wraps
 
 from bson.son import SON
 from pymongo import AsyncMongoClient
@@ -154,6 +156,42 @@ 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):
+    is_cpython_linux = sys.platform == "linux" and sys.implementation.name == "cpython"
+    if is_cpython_linux and not affects_cpython_linux:
+        max_runs = 1
+        min_passes = 1
+
+    def decorator(target_func):
+        @wraps(target_func)
+        async def wrapper(*args, **kwargs):
+            passes = 0
+            failure = None
+            for i in range(max_runs):
+                try:
+                    result = await target_func(*args, **kwargs)
+                    passes += 1
+                    if passes == min_passes:
+                        return result
+                except Exception as e:
+                    failure = e
+                    await asyncio.sleep(delay)
+                if failure is not None and i < max_runs - 1:
+                    print("Flaky failure:", failure)
+            if failure:
+                raise failure
+            raise RuntimeError(f"Only passed {passes} of {min_passes} times")
+
+        return wrapper
+
+    # If `func` is callable, the decorator was used without arguments (`@flaky`)
+    if callable(func):
+        return decorator(func)
+
+    # Otherwise, return the decorator function, allowing arguments (`@flaky(max_runs=...)`)
+    return decorator
+
+
 class AsyncMockConnection:
     def __init__(self):
         self.cancel_context = _CancellationContext()

test/test_client_bulk_write.py

@@ -25,6 +25,7 @@
     client_context,
     unittest,
 )
+from test.utils import flaky
 from test.utils_shared import (
     OvertCommandListener,
 )
@@ -615,15 +616,14 @@ def test_15_unacknowledged_write_across_batches(self):
 # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#11-multi-batch-bulkwrites
 class TestClientBulkWriteCSOT(IntegrationTest):
     def setUp(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         super().setUp()
         self.max_write_batch_size = client_context.max_write_batch_size
         self.max_bson_object_size = client_context.max_bson_size
         self.max_message_size_bytes = client_context.max_message_size_bytes
 
     @client_context.require_version_min(8, 0, 0, -24)
     @client_context.require_failCommand_fail_point
+    @flaky
     def test_timeout_in_multi_batch_bulk_write(self):
         _OVERHEAD = 500
 

test/test_csot.py

@@ -23,6 +23,7 @@
 
 from test import IntegrationTest, client_context, unittest
 from test.unified_format import generate_test_classes
+from test.utils import flaky
 
 import pymongo
 from pymongo import _csot
@@ -43,9 +44,8 @@
 class TestCSOT(IntegrationTest):
     RUN_ON_LOAD_BALANCER = True
 
+    @flaky
     def test_timeout_nested(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         coll = self.db.coll
         self.assertEqual(_csot.get_timeout(), None)
         self.assertEqual(_csot.get_deadline(), float("inf"))
@@ -82,9 +82,8 @@ def test_timeout_nested(self):
         self.assertEqual(_csot.get_rtt(), 0.0)
 
     @client_context.require_change_streams
+    @flaky
     def test_change_stream_can_resume_after_timeouts(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         coll = self.db.test
         coll.insert_one({})
         with coll.watch() as stream:

test/test_cursor.py

@@ -31,6 +31,7 @@
 sys.path[0:0] = [""]
 
 from test import IntegrationTest, client_context, unittest
+from test.utils import flaky
 from test.utils_shared import (
     AllowListEventListener,
     EventListener,
@@ -1406,9 +1407,8 @@ def test_to_list_length(self):
         docs = c.to_list(3)
         self.assertEqual(len(docs), 2)
 
+    @flaky
     def test_to_list_csot_applied(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         client = self.single_client(timeoutMS=500, w=1)
         coll = client.pymongo.test
         # Initialize the client with a larger timeout to help make test less flakey
@@ -1449,9 +1449,8 @@ def test_command_cursor_to_list_length(self):
         self.assertEqual(len(result.to_list(1)), 1)
 
     @client_context.require_failCommand_blockConnection
+    @flaky
     def test_command_cursor_to_list_csot_applied(self):
-        if os.environ.get("SKIP_CSOT_TESTS", ""):
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
         client = self.single_client(timeoutMS=500, w=1)
         coll = client.pymongo.test
         # Initialize the client with a larger timeout to help make test less flakey

test/unified_format.py

@@ -38,7 +38,6 @@
 from test.unified_format_shared import (
     KMS_TLS_OPTS,
     PLACEHOLDER_MAP,
-    SKIP_CSOT_TESTS,
     EventListenerUtil,
     MatchEvaluatorUtil,
     coerce_result,
@@ -1361,9 +1360,6 @@ def verify_outcome(self, spec):
                 self.assertListEqual(sorted_expected_documents, actual_documents)
 
     def run_scenario(self, spec, uri=None):
-        if "csot" in self.id().lower() and SKIP_CSOT_TESTS:
-            raise unittest.SkipTest("SKIP_CSOT_TESTS is set, skipping...")
-
         # Kill all sessions before and after each test to prevent an open
         # transaction (from a test failure) from blocking collection/database
         # operations during test set up and tear down.