PYTHON-5413 Handle flaky tests (#2395)

Author: blink1073

.evergreen/generated_configs/functions.yml

@@ -145,6 +145,7 @@ functions:
           - MONGODB_API_VERSION
           - REQUIRE_API_VERSION
           - DEBUG_LOG
+          - DISABLE_FLAKY
           - ORCHESTRATION_FILE
           - OCSP_SERVER_TYPE
           - VERSION

.evergreen/scripts/generate_config.py

@@ -1084,6 +1084,7 @@ def create_run_tests_func():
         "MONGODB_API_VERSION",
         "REQUIRE_API_VERSION",
         "DEBUG_LOG",
+        "DISABLE_FLAKY",
         "ORCHESTRATION_FILE",
         "OCSP_SERVER_TYPE",
         "VERSION",

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

CONTRIBUTING.md

@@ -404,6 +404,15 @@ If you are running one of the `no-responder` tests, omit the `run-server` step.
 - Regenerate the test variants and tasks using `pre-commit run --all-files generate-config`.
 - Make sure to add instructions for running the test suite to `CONTRIBUTING.md`.
 
+## Handling flaky tests
+
+We have a custom `flaky` decorator in [test/asynchronous/utils.py](test/asynchronous/utils.py) that can be used for
+tests that are `flaky`.  By default the decorator only applies when not running on CPython on Linux, since other
+runtimes tend to have more variation.  When using the `flaky` decorator, open a corresponding ticket and
+a use the ticket number as the "reason" parameter to the decorator, e.g. `@flaky(reason="PYTHON-1234")`.
+When running tests locally (not in CI), the `flaky` decorator will be disabled unless `ENABLE_FLAKY` is set.
+To disable the `flaky` decorator in CI, you can use `evergreen patch --param DISABLE_FLAKY=1`.
+
 ## Specification Tests
 
 The MongoDB [specifications repository](https://github.com/mongodb/specifications)

test/__init__.py

@@ -32,6 +32,7 @@
 import warnings
 from asyncio import iscoroutinefunction
 
+from pymongo.errors import AutoReconnect
 from pymongo.synchronous.uri_parser import parse_uri
 
 try:
@@ -1219,12 +1220,17 @@ def teardown():
     c = client_context.client
     if c:
         if not client_context.is_data_lake:
-            c.drop_database("pymongo-pooling-tests")
-            c.drop_database("pymongo_test")
-            c.drop_database("pymongo_test1")
-            c.drop_database("pymongo_test2")
-            c.drop_database("pymongo_test_mike")
-            c.drop_database("pymongo_test_bernie")
+            try:
+                c.drop_database("pymongo-pooling-tests")
+                c.drop_database("pymongo_test")
+                c.drop_database("pymongo_test1")
+                c.drop_database("pymongo_test2")
+                c.drop_database("pymongo_test_mike")
+                c.drop_database("pymongo_test_bernie")
+            except AutoReconnect:
+                # PYTHON-4982
+                if sys.implementation.name.lower() != "pypy":
+                    raise
         c.close()
     print_running_clients()
 

test/asynchronous/__init__.py

@@ -33,6 +33,7 @@
 from asyncio import iscoroutinefunction
 
 from pymongo.asynchronous.uri_parser import parse_uri
+from pymongo.errors import AutoReconnect
 
 try:
     import ipaddress
@@ -1235,12 +1236,17 @@ async def async_teardown():
     c = async_client_context.client
     if c:
         if not async_client_context.is_data_lake:
-            await c.drop_database("pymongo-pooling-tests")
-            await c.drop_database("pymongo_test")
-            await c.drop_database("pymongo_test1")
-            await c.drop_database("pymongo_test2")
-            await c.drop_database("pymongo_test_mike")
-            await c.drop_database("pymongo_test_bernie")
+            try:
+                await c.drop_database("pymongo-pooling-tests")
+                await c.drop_database("pymongo_test")
+                await c.drop_database("pymongo_test1")
+                await c.drop_database("pymongo_test2")
+                await c.drop_database("pymongo_test_mike")
+                await c.drop_database("pymongo_test_bernie")
+            except AutoReconnect:
+                # PYTHON-4982
+                if sys.implementation.name.lower() != "pypy":
+                    raise
         await c.close()
     print_running_clients()
 

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,16 +620,17 @@ 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(reason="PYTHON-5290", max_runs=3, affects_cpython_linux=True)
     async def test_timeout_in_multi_batch_bulk_write(self):
+        if sys.platform != "linux" and "CI" in os.environ:
+            self.skipTest("PYTHON-3522 CSOT test runs too slow on Windows and MacOS")
         _OVERHEAD = 500
 
         internal_client = await self.async_rs_or_single_client(timeoutMS=None)

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(reason="PYTHON-3522")
     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(reason="PYTHON-3522")
     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,
@@ -1406,9 +1407,8 @@ async def test_to_list_length(self):
         docs = await c.to_list(3)
         self.assertEqual(len(docs), 2)
 
+    @flaky(reason="PYTHON-3522")
     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
@@ -1449,9 +1449,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(reason="PYTHON-3522")
     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/test_encryption.py

@@ -32,6 +32,7 @@
 import warnings
 from test.asynchronous import AsyncIntegrationTest, AsyncPyMongoTestCase, async_client_context
 from test.asynchronous.test_bulk import AsyncBulkTestBase
+from test.asynchronous.utils import flaky
 from test.asynchronous.utils_spec_runner import AsyncSpecRunner, AsyncSpecTestCreator
 from threading import Thread
 from typing import Any, Dict, Mapping, Optional
@@ -3247,6 +3248,7 @@ async def test_kms_retry(self):
 class TestAutomaticDecryptionKeys(AsyncEncryptionIntegrationTest):
     @async_client_context.require_no_standalone
     @async_client_context.require_version_min(7, 0, -1)
+    @flaky(reason="PYTHON-4982")
     async def asyncSetUp(self):
         await super().asyncSetUp()
         self.key1_document = json_data("etc", "data", "keys", "key1-document.json")
@@ -3489,6 +3491,8 @@ async def test_implicit_session_ignored_when_unsupported(self):
 
         self.assertNotIn("lsid", self.listener.started_events[1].command)
 
+        await self.mongocryptd_client.close()
+
     async def test_explicit_session_errors_when_unsupported(self):
         self.listener.reset()
         async with self.mongocryptd_client.start_session() as s:
@@ -3501,6 +3505,8 @@ async def test_explicit_session_errors_when_unsupported(self):
             ):
                 await self.mongocryptd_client.db.test.insert_one({"x": 1}, session=s)
 
+        await self.mongocryptd_client.close()
+
 
 if __name__ == "__main__":
     unittest.main()