PYTHON-5474 Wait until all mongoses are discovered for mongos retry tests

test/asynchronous/test_retryable_reads.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import os
 import pprint
 import sys
 import threading
@@ -25,7 +24,11 @@
 from pymongo import MongoClient
 from pymongo.common import MAX_ADAPTIVE_RETRIES
 from pymongo.errors import OperationFailure, PyMongoError
-from test.asynchronous.utils import async_ensure_all_connected, async_set_fail_point
+from test.asynchronous.utils import (
+    async_ensure_all_connected,
+    async_set_fail_point,
+    async_wait_until,
+)
 
 sys.path[0:0] = [""]
 
@@ -182,6 +185,14 @@ async def test_retryable_reads_are_retried_on_a_different_mongos_when_one_is_ava
             retryReads=True,
         )
 
+        # Ensure the client has discovered all mongoses.
+        await async_wait_until(
+            lambda: (
+                len(client.topology_description.known_servers) == len(async_client_context.mongoses)
+            ),
+            "discover all mongoses",
+        )
+
         with self.assertRaises(OperationFailure):
             await client.db.coll.find_one({})
 

test/asynchronous/test_retryable_writes.py

@@ -24,7 +24,7 @@
 from unittest import mock
 
 from pymongo.common import MAX_ADAPTIVE_RETRIES
-from test.asynchronous.utils import async_set_fail_point, flaky
+from test.asynchronous.utils import async_set_fail_point, async_wait_until, flaky
 
 sys.path[0:0] = [""]
 
@@ -395,6 +395,14 @@ async def test_retryable_writes_in_sharded_cluster_multiple_available(self):
             retryWrites=True,
         )
 
+        # Ensure the client has discovered all mongoses.
+        await async_wait_until(
+            lambda: (
+                len(client.topology_description.known_servers) == len(async_client_context.mongoses)
+            ),
+            "discover all mongoses",
+        )
+
         with self.assertRaises(AutoReconnect):
             await client.db.coll.insert_one({"x": 1})
 

test/test_retryable_reads.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import os
 import pprint
 import sys
 import threading
@@ -25,7 +24,11 @@
 from pymongo import MongoClient
 from pymongo.common import MAX_ADAPTIVE_RETRIES
 from pymongo.errors import OperationFailure, PyMongoError
-from test.utils import ensure_all_connected, set_fail_point
+from test.utils import (
+    ensure_all_connected,
+    set_fail_point,
+    wait_until,
+)
 
 sys.path[0:0] = [""]
 
@@ -182,6 +185,12 @@ def test_retryable_reads_are_retried_on_a_different_mongos_when_one_is_available
             retryReads=True,
         )
 
+        # Ensure the client has discovered all mongoses.
+        wait_until(
+            lambda: len(client.topology_description.known_servers) == len(client_context.mongoses),
+            "discover all mongoses",
+        )
+
         with self.assertRaises(OperationFailure):
             client.db.coll.find_one({})
 

test/test_retryable_writes.py

@@ -24,7 +24,7 @@
 from unittest import mock
 
 from pymongo.common import MAX_ADAPTIVE_RETRIES
-from test.utils import flaky, set_fail_point
+from test.utils import flaky, set_fail_point, wait_until
 
 sys.path[0:0] = [""]
 
@@ -393,6 +393,12 @@ def test_retryable_writes_in_sharded_cluster_multiple_available(self):
             retryWrites=True,
         )
 
+        # Ensure the client has discovered all mongoses.
+        wait_until(
+            lambda: len(client.topology_description.known_servers) == len(client_context.mongoses),
+            "discover all mongoses",
+        )
+
         with self.assertRaises(AutoReconnect):
             client.db.coll.insert_one({"x": 1})