make an async ensure all connected

Author: sleepyStick

test/asynchronous/test_connections_survive_primary_stepdown_spec.py

@@ -22,7 +22,7 @@
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
 from test.utils import (
     CMAPListener,
-    ensure_all_connected,
+    async_ensure_all_connected,
     repl_set_step_down,
 )
 
@@ -51,7 +51,7 @@ async def _setup_class(cls):
         # Ensure connections to all servers in replica set. This is to test
         # that the is_writable flag is properly updated for connections that
         # survive a replica set election.
-        ensure_all_connected(cls.client)
+        async_ensure_all_connected(cls.client)
         cls.listener.reset()
 
         cls.db = cls.client.get_database("step-down", write_concern=WriteConcern("majority"))

test/test_connections_survive_primary_stepdown_spec.py

@@ -22,7 +22,7 @@
 from test import IntegrationTest, client_context, unittest
 from test.utils import (
     CMAPListener,
-    ensure_all_connected,
+    async_ensure_all_connected,
     repl_set_step_down,
 )
 
@@ -51,7 +51,7 @@ def _setup_class(cls):
         # Ensure connections to all servers in replica set. This is to test
         # that the is_writable flag is properly updated for connections that
         # survive a replica set election.
-        ensure_all_connected(cls.client)
+        async_ensure_all_connected(cls.client)
         cls.listener.reset()
 
         cls.db = cls.client.get_database("step-down", write_concern=WriteConcern("majority"))

test/utils.py

@@ -599,6 +599,40 @@ def discover():
         )
 
 
+async def async_ensure_all_connected(client: AsyncMongoClient) -> None:
+    """Ensure that the client's connection pool has socket connections to all
+    members of a replica set. Raises ConfigurationError when called with a
+    non-replica set client.
+
+    Depending on the use-case, the caller may need to clear any event listeners
+    that are configured on the client.
+    """
+    hello: dict = await client.admin.command(HelloCompat.LEGACY_CMD)
+    if "setName" not in hello:
+        raise ConfigurationError("cluster is not a replica set")
+
+    target_host_list = set(hello["hosts"] + hello.get("passives", []))
+    connected_host_list = {hello["me"]}
+
+    # Run hello until we have connected to each host at least once.
+    def discover():
+        i = 0
+        while i < 100 and connected_host_list != target_host_list:
+            hello: dict = client.admin.command(
+                HelloCompat.LEGACY_CMD, read_preference=ReadPreference.SECONDARY
+            )
+            connected_host_list.update([hello["me"]])
+            i += 1
+        return connected_host_list
+
+    try:
+        async_wait_until(lambda: target_host_list == discover(), "connected to all hosts")
+    except AssertionError as exc:
+        raise AssertionError(
+            f"{exc}, {connected_host_list} != {target_host_list}, {client.topology_description}"
+        )
+
+
 def one(s):
     """Get one element of a set"""
     return next(iter(s))