From 8ba2feb8b83abf349f3a252ca6fe07ee41f2c577 Mon Sep 17 00:00:00 2001
From: Shane Harvey <shnhrv@gmail.com>
Date: Mon, 9 Jun 2025 11:13:00 -0700
Subject: [PATCH] PYTHON-5409 Make test_implicit_sessions_checkout less flaky

---
 test/asynchronous/test_session.py | 13 +++++++------
 test/test_session.py              | 13 +++++++------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/test/asynchronous/test_session.py b/test/asynchronous/test_session.py
index 1f6fb0d319..0ceaea98f9 100644
--- a/test/asynchronous/test_session.py
+++ b/test/asynchronous/test_session.py
@@ -194,10 +194,11 @@ async def test_implicit_sessions_checkout(self):
         # successful connection checkout" test from Driver Sessions Spec.
         succeeded = False
         lsid_set = set()
-        failures = 0
-        for _ in range(5):
-            listener = OvertCommandListener()
-            client = await self.async_rs_or_single_client(event_listeners=[listener], maxPoolSize=1)
+        listener = OvertCommandListener()
+        client = await self.async_rs_or_single_client(event_listeners=[listener], maxPoolSize=1)
+        # Retry up to 10 times because there is a known race that can cause multiple
+        # sessions to be used: connection check in happens before session check in
+        for _ in range(10):
             cursor = client.db.test.find({})
             ops: List[Tuple[Callable, List[Any]]] = [
                 (client.db.test.find_one, [{"_id": 1}]),
@@ -240,9 +241,9 @@ async def target(op, *args):
                 if i.command.get("lsid"):
                     lsid_set.add(i.command.get("lsid")["id"])
             if len(lsid_set) == 1:
+                # Break on first success.
                 succeeded = True
-            else:
-                failures += 1
+                break
         self.assertTrue(succeeded, lsid_set)
 
     async def test_pool_lifo(self):
diff --git a/test/test_session.py b/test/test_session.py
index 49cb9dba91..d70032d15f 100644
--- a/test/test_session.py
+++ b/test/test_session.py
@@ -194,10 +194,11 @@ def test_implicit_sessions_checkout(self):
         # successful connection checkout" test from Driver Sessions Spec.
         succeeded = False
         lsid_set = set()
-        failures = 0
-        for _ in range(5):
-            listener = OvertCommandListener()
-            client = self.rs_or_single_client(event_listeners=[listener], maxPoolSize=1)
+        listener = OvertCommandListener()
+        client = self.rs_or_single_client(event_listeners=[listener], maxPoolSize=1)
+        # Retry up to 10 times because there is a known race that can cause multiple
+        # sessions to be used: connection check in happens before session check in
+        for _ in range(10):
             cursor = client.db.test.find({})
             ops: List[Tuple[Callable, List[Any]]] = [
                 (client.db.test.find_one, [{"_id": 1}]),
@@ -240,9 +241,9 @@ def target(op, *args):
                 if i.command.get("lsid"):
                     lsid_set.add(i.command.get("lsid")["id"])
             if len(lsid_set) == 1:
+                # Break on first success.
                 succeeded = True
-            else:
-                failures += 1
+                break
         self.assertTrue(succeeded, lsid_set)
 
     def test_pool_lifo(self):