PYTHON-5356 - Init unified test clients for CSOT tests

Author: NoahStapp

test/asynchronous/unified_format.py

@@ -257,7 +257,7 @@ def _handle_placeholders(self, spec: dict, current: dict, path: str) -> Any:
                 current[key] = self._handle_placeholders(spec, value, subpath)
         return current
 
-    async def _create_entity(self, entity_spec, uri=None):
+    async def _create_entity(self, entity_spec, uri=None, init_client=False):
         if len(entity_spec) != 1:
             self.test.fail(f"Entity spec {entity_spec} did not contain exactly one top-level key")
 
@@ -303,6 +303,8 @@ async def _create_entity(self, entity_spec, uri=None):
             if uri:
                 kwargs["h"] = uri
             client = await self.test.async_rs_or_single_client(**kwargs)
+            if init_client:
+                await client.aconnect()
             self[spec["id"]] = client
             return
         elif entity_type == "database":
@@ -390,9 +392,9 @@ async def drop(self: AsyncGridFSBucket, *args: Any, **kwargs: Any) -> None:
 
         self.test.fail(f"Unable to create entity of unknown type {entity_type}")
 
-    async def create_entities_from_spec(self, entity_spec, uri=None):
+    async def create_entities_from_spec(self, entity_spec, uri=None, init_client=False):
         for spec in entity_spec:
-            await self._create_entity(spec, uri=uri)
+            await self._create_entity(spec, uri=uri, init_client=init_client)
 
     def get_listener_for_client(self, client_name: str) -> EventListenerUtil:
         client = self[client_name]
@@ -1406,7 +1408,7 @@ async def run_scenario(self, spec, uri=None):
             attempts = 3
             for i in range(attempts):
                 try:
-                    return await self._run_scenario(spec, uri)
+                    return await self._run_scenario(spec, uri, init_client=True)
                 except (AssertionError, OperationFailure) as exc:
                     if isinstance(exc, OperationFailure) and (
                         _IS_SYNC or "failpoint" not in exc._message
@@ -1426,7 +1428,7 @@ async def run_scenario(self, spec, uri=None):
             await self._run_scenario(spec, uri)
             return None
 
-    async def _run_scenario(self, spec, uri=None):
+    async def _run_scenario(self, spec, uri=None, init_client=False):
         # maybe skip test manually
         self.maybe_skip_test(spec)
 
@@ -1444,7 +1446,7 @@ async def _run_scenario(self, spec, uri=None):
         self._uri = uri
         self.entity_map = EntityMapUtil(self)
         await self.entity_map.create_entities_from_spec(
-            self.TEST_SPEC.get("createEntities", []), uri=uri
+            self.TEST_SPEC.get("createEntities", []), uri=uri, init_client=init_client
         )
         self._cluster_time = None
         # process initialData

test/unified_format.py

@@ -256,7 +256,7 @@ def _handle_placeholders(self, spec: dict, current: dict, path: str) -> Any:
                 current[key] = self._handle_placeholders(spec, value, subpath)
         return current
 
-    def _create_entity(self, entity_spec, uri=None):
+    def _create_entity(self, entity_spec, uri=None, init_client=False):
         if len(entity_spec) != 1:
             self.test.fail(f"Entity spec {entity_spec} did not contain exactly one top-level key")
 
@@ -302,6 +302,8 @@ def _create_entity(self, entity_spec, uri=None):
             if uri:
                 kwargs["h"] = uri
             client = self.test.rs_or_single_client(**kwargs)
+            if init_client:
+                client._connect()
             self[spec["id"]] = client
             return
         elif entity_type == "database":
@@ -389,9 +391,9 @@ def drop(self: GridFSBucket, *args: Any, **kwargs: Any) -> None:
 
         self.test.fail(f"Unable to create entity of unknown type {entity_type}")
 
-    def create_entities_from_spec(self, entity_spec, uri=None):
+    def create_entities_from_spec(self, entity_spec, uri=None, init_client=False):
         for spec in entity_spec:
-            self._create_entity(spec, uri=uri)
+            self._create_entity(spec, uri=uri, init_client=init_client)
 
     def get_listener_for_client(self, client_name: str) -> EventListenerUtil:
         client = self[client_name]
@@ -1393,7 +1395,7 @@ def run_scenario(self, spec, uri=None):
             attempts = 3
             for i in range(attempts):
                 try:
-                    return self._run_scenario(spec, uri)
+                    return self._run_scenario(spec, uri, init_client=True)
                 except (AssertionError, OperationFailure) as exc:
                     if isinstance(exc, OperationFailure) and (
                         _IS_SYNC or "failpoint" not in exc._message
@@ -1413,7 +1415,7 @@ def run_scenario(self, spec, uri=None):
             self._run_scenario(spec, uri)
             return None
 
-    def _run_scenario(self, spec, uri=None):
+    def _run_scenario(self, spec, uri=None, init_client=False):
         # maybe skip test manually
         self.maybe_skip_test(spec)
 
@@ -1430,7 +1432,9 @@ def _run_scenario(self, spec, uri=None):
         # process createEntities
         self._uri = uri
         self.entity_map = EntityMapUtil(self)
-        self.entity_map.create_entities_from_spec(self.TEST_SPEC.get("createEntities", []), uri=uri)
+        self.entity_map.create_entities_from_spec(
+            self.TEST_SPEC.get("createEntities", []), uri=uri, init_client=init_client
+        )
         self._cluster_time = None
         # process initialData
         if "initialData" in self.TEST_SPEC: