PYTHON-3405/PYTHON-2531 Fix tests for primary step down (#1041)

Author: ShaneHarvey

test/unified_format.py

@@ -114,6 +114,7 @@
 from pymongo.server_selectors import Selection, writable_server_selector
 from pymongo.server_type import SERVER_TYPE
 from pymongo.topology_description import TopologyDescription
+from pymongo.typings import _Address
 from pymongo.write_concern import WriteConcern
 
 JSON_OPTS = json_util.JSONOptions(tz_aware=False)
@@ -1442,21 +1443,21 @@ def _testOperation_assertTopologyType(self, spec):
         self.assertIsInstance(description, TopologyDescription)
         self.assertEqual(description.topology_type_name, spec["topologyType"])
 
-    def _testOperation_waitForPrimaryChange(self, spec):
+    def _testOperation_waitForPrimaryChange(self, spec: dict) -> None:
         """Run the waitForPrimaryChange test operation."""
         client = self.entity_map[spec["client"]]
         old_description: TopologyDescription = self.entity_map[spec["priorTopologyDescription"]]
         timeout = spec["timeoutMS"] / 1000.0
 
-        def get_primary(td: TopologyDescription) -> Optional[ServerDescription]:
+        def get_primary(td: TopologyDescription) -> Optional[_Address]:
             servers = writable_server_selector(Selection.from_topology_description(td))
             if servers and servers[0].server_type == SERVER_TYPE.RSPrimary:
-                return servers[0]
+                return servers[0].address
             return None
 
         old_primary = get_primary(old_description)
 
-        def primary_changed():
+        def primary_changed() -> bool:
             primary = client.primary
             if primary is None:
                 return False

test/utils.py

@@ -593,7 +593,7 @@ def rs_or_single_client(h=None, p=None, **kwargs):
     return _mongo_client(h, p, **kwargs)
 
 
-def ensure_all_connected(client):
+def ensure_all_connected(client: MongoClient) -> 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.
@@ -605,14 +605,26 @@ def ensure_all_connected(client):
     if "setName" not in hello:
         raise ConfigurationError("cluster is not a replica set")
 
-    target_host_list = set(hello["hosts"])
+    target_host_list = set(hello["hosts"] + hello.get("passives", []))
     connected_host_list = set([hello["me"]])
-    admindb = client.get_database("admin")
 
     # Run hello until we have connected to each host at least once.
-    while connected_host_list != target_host_list:
-        hello = admindb.command(HelloCompat.LEGACY_CMD, read_preference=ReadPreference.SECONDARY)
-        connected_host_list.update([hello["me"]])
+    def discover():
+        i = 0
+        while i < 100 and connected_host_list != target_host_list:
+            hello = client.admin.command(
+                HelloCompat.LEGACY_CMD, read_preference=ReadPreference.SECONDARY
+            )
+            connected_host_list.update([hello["me"]])
+            i += 1
+        return connected_host_list
+
+    try:
+        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):