PYTHON-5089 Convert test.test_mongos_load_balancing to async

test/asynchronous/test_mongos_load_balancing.py

@@ -34,17 +34,11 @@
 _IS_SYNC = False
 
 
-@async_client_context.require_connection
-@async_client_context.require_no_load_balancer
-def asyncSetUpModule():
-    pass
-
-
 if not _IS_SYNC:
 
     class SimpleOp:
         def __init__(self, client):
-            self.task: asyncio.Task
+            self.task = asyncio.create_task(self.run())
             self.client = client
             self.passed = False
 
@@ -53,7 +47,7 @@ async def run(self):
             self.passed = True  # No exception raised.
 
         def start(self):
-            self.task = asyncio.create_task(self.run())
+            pass
 
         async def join(self):
             await self.task
@@ -70,15 +64,15 @@ def run(self):
             self.passed = True  # No exception raised.
 
 
-async def do_simple_op(client, nthreads):
-    threads = [SimpleOp(client) for _ in range(nthreads)]
-    for t in threads:
+async def do_simple_op(client, ntasks):
+    tasks = [SimpleOp(client) for _ in range(ntasks)]
+    for t in tasks:
         t.start()
 
-    for t in threads:
+    for t in tasks:
         await t.join()
 
-    for t in threads:
+    for t in tasks:
         assert t.passed
 
 
@@ -90,6 +84,11 @@ async def writable_addresses(topology):
 
 
 class TestMongosLoadBalancing(AsyncMockClientTest):
+    @async_client_context.require_connection
+    @async_client_context.require_no_load_balancer
+    async def asyncSetUp(self):
+        await super().asyncSetUp()
+
     def mock_client(self, **kwargs):
         mock_client = AsyncMockClient(
             standalones=[],

test/test_mongos_load_balancing.py

@@ -34,17 +34,11 @@
 _IS_SYNC = True
 
 
-@client_context.require_connection
-@client_context.require_no_load_balancer
-def setUpModule():
-    pass
-
-
 if not _IS_SYNC:
 
     class SimpleOp:
         def __init__(self, client):
-            self.task: asyncio.Task
+            self.task = asyncio.create_task(self.run())
             self.client = client
             self.passed = False
 
@@ -53,7 +47,7 @@ def run(self):
             self.passed = True  # No exception raised.
 
         def start(self):
-            self.task = asyncio.create_task(self.run())
+            pass
 
         def join(self):
             self.task
@@ -70,15 +64,15 @@ def run(self):
             self.passed = True  # No exception raised.
 
 
-def do_simple_op(client, nthreads):
-    threads = [SimpleOp(client) for _ in range(nthreads)]
-    for t in threads:
+def do_simple_op(client, ntasks):
+    tasks = [SimpleOp(client) for _ in range(ntasks)]
+    for t in tasks:
         t.start()
 
-    for t in threads:
+    for t in tasks:
         t.join()
 
-    for t in threads:
+    for t in tasks:
         assert t.passed
 
 
@@ -90,6 +84,11 @@ def writable_addresses(topology):
 
 
 class TestMongosLoadBalancing(MockClientTest):
+    @client_context.require_connection
+    @client_context.require_no_load_balancer
+    def setUp(self):
+        super().setUp()
+
     def mock_client(self, **kwargs):
         mock_client = MockClient(
             standalones=[],