make test_heartbeat_start_ordering async

Author: sleepyStick

test/asynchronous/test_discovery_and_monitoring.py

@@ -20,7 +20,7 @@
 import socketserver
 import sys
 import threading
-from asyncio import StreamReader
+from asyncio import StreamReader, StreamWriter
 from pathlib import Path
 
 sys.path[0:0] = [""]
@@ -227,7 +227,7 @@ async def run_scenario(self):
     return run_scenario
 
 
-async def create_tests():
+def create_tests():
     for dirpath, _, filenames in os.walk(SDAM_PATH):
         dirname = os.path.split(dirpath)[-1]
         # SDAM unified tests are handled separately.
@@ -248,8 +248,6 @@ async def create_tests():
             setattr(TestAllScenarios, new_test.__name__, new_test)
 
 
-
-
 class TestClusterTimeComparison(unittest.IsolatedAsyncioTestCase):
     async def test_cluster_time_comparison(self):
         t = await create_mock_topology("mongodb://host")
@@ -278,6 +276,7 @@ async def send_cluster_time(time, inc, should_update):
 
 class TestIgnoreStaleErrors(AsyncIntegrationTest):
     if _IS_SYNC:
+
         async def test_ignore_stale_connection_errors(self):
             N_THREADS = 5
             barrier = threading.Barrier(N_THREADS, timeout=30)
@@ -317,6 +316,7 @@ async def insert_command(i):
             # Server should be selectable.
             await client.admin.command("ping")
     else:
+
         async def test_ignore_stale_connection_errors(self):
             N_TASKS = 5
             barrier = asyncio.Barrier(N_TASKS)
@@ -469,62 +469,61 @@ def handle_request_and_shutdown(self):
 
 
 class TestHeartbeatStartOrdering(AsyncPyMongoTestCase):
-    if _IS_SYNC:
-        async def test_heartbeat_start_ordering(self):
-            events = []
-            listener = HeartbeatEventsListListener(events)
+    async def test_heartbeat_start_ordering(self):
+        events = []
+        listener = HeartbeatEventsListListener(events)
+
+        if _IS_SYNC:
             server = TCPServer(("localhost", 9999), MockTCPHandler)
             server.events = events
             server_thread = threading.Thread(target=server.handle_request_and_shutdown)
             server_thread.start()
             _c = await self.simple_client(
-                "mongodb://localhost:9999", serverSelectionTimeoutMS=500, event_listeners=(listener,)
+                "mongodb://localhost:9999",
+                serverSelectionTimeoutMS=500,
+                event_listeners=(listener,),
             )
             server_thread.join()
             listener.wait_for_event(ServerHeartbeatStartedEvent, 1)
             listener.wait_for_event(ServerHeartbeatFailedEvent, 1)
 
-            self.assertEqual(
-                events,
-                [
-                    "serverHeartbeatStartedEvent",
-                    "client connected",
-                    "client hello received",
-                    "serverHeartbeatFailedEvent",
-                ],
-            )
-    else:
-        async def test_heartbeat_start_ordering(self):
-            events = []
+        else:
 
-            async def handle_client(reader: StreamReader, writer):
-                server.events.append("client connected")
-                print("clent connected")
+            async def handle_client(reader: StreamReader, writer: StreamWriter):
+                events.append("client connected")
                 if (await reader.read(1024)).strip():
-                    server.events.append("client hello received")
-                    print("client helllo recieved")
-            listener = HeartbeatEventsListListener(events)
+                    events.append("client hello received")
+                writer.close()
+                await writer.wait_closed()
+
             server = await asyncio.start_server(handle_client, "localhost", 9999)
-            async with server:
-                server.events = events
-                _c = self.simple_client(
-                    "mongodb://localhost:9999", serverSelectionTimeoutMS=500, event_listeners=(listener,)
-                )
-                server.close()
-                server_task = asyncio.create_task(server.wait_closed())
-                await server_task
-                await listener.async_wait_for_event(ServerHeartbeatStartedEvent, 1)
-                await listener.async_wait_for_event(ServerHeartbeatFailedEvent, 1)
-
-                self.assertEqual(
-                    events,
-                    [
-                        "serverHeartbeatStartedEvent",
-                        "client connected",
-                        "client hello received",
-                        "serverHeartbeatFailedEvent",
-                    ],
-                )
+            server.events = events
+            await server.start_serving()
+            print(server.is_serving())
+            _c = self.simple_client(
+                "mongodb://localhost:9999",
+                serverSelectionTimeoutMS=500,
+                event_listeners=(listener,),
+            )
+            if _c._options.connect:
+                await _c.aconnect()
+
+            await listener.async_wait_for_event(ServerHeartbeatStartedEvent, 1)
+            await listener.async_wait_for_event(ServerHeartbeatFailedEvent, 1)
+
+            server.close()
+            await server.wait_closed()
+            await _c.close()
+
+        self.assertEqual(
+            events,
+            [
+                "serverHeartbeatStartedEvent",
+                "client connected",
+                "client hello received",
+                "serverHeartbeatFailedEvent",
+            ],
+        )
 
 
 # Generate unified tests.

test/asynchronous/unified_format.py

@@ -1155,7 +1155,7 @@ def _testOperation_assertTopologyType(self, spec):
         self.assertIsInstance(description, TopologyDescription)
         self.assertEqual(description.topology_type_name, spec["topologyType"])
 
-    def _testOperation_waitForPrimaryChange(self, spec: dict) -> None:
+    async 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"]]
@@ -1169,13 +1169,13 @@ def get_primary(td: TopologyDescription) -> Optional[_Address]:
 
         old_primary = get_primary(old_description)
 
-        def primary_changed() -> bool:
-            primary = client.primary
+        async def primary_changed() -> bool:
+            primary = await client.primary
             if primary is None:
                 return False
             return primary != old_primary
 
-        wait_until(primary_changed, "change primary", timeout=timeout)
+        await async_wait_until(primary_changed, "change primary", timeout=timeout)
 
     async def _testOperation_runOnThread(self, spec):
         """Run the 'runOnThread' operation."""

test/test_discovery_and_monitoring.py

@@ -20,7 +20,7 @@
 import socketserver
 import sys
 import threading
-from asyncio import StreamReader
+from asyncio import StreamReader, StreamWriter
 from pathlib import Path
 
 sys.path[0:0] = [""]
@@ -248,8 +248,6 @@ def create_tests():
             setattr(TestAllScenarios, new_test.__name__, new_test)
 
 
-
-
 class TestClusterTimeComparison(unittest.TestCase):
     def test_cluster_time_comparison(self):
         t = create_mock_topology("mongodb://host")
@@ -278,6 +276,7 @@ def send_cluster_time(time, inc, should_update):
 
 class TestIgnoreStaleErrors(IntegrationTest):
     if _IS_SYNC:
+
         def test_ignore_stale_connection_errors(self):
             N_THREADS = 5
             barrier = threading.Barrier(N_THREADS, timeout=30)
@@ -317,6 +316,7 @@ def insert_command(i):
             # Server should be selectable.
             client.admin.command("ping")
     else:
+
         def test_ignore_stale_connection_errors(self):
             N_TASKS = 5
             barrier = asyncio.Barrier(N_TASKS)
@@ -469,62 +469,61 @@ def handle_request_and_shutdown(self):
 
 
 class TestHeartbeatStartOrdering(PyMongoTestCase):
-    if _IS_SYNC:
-        def test_heartbeat_start_ordering(self):
-            events = []
-            listener = HeartbeatEventsListListener(events)
+    def test_heartbeat_start_ordering(self):
+        events = []
+        listener = HeartbeatEventsListListener(events)
+
+        if _IS_SYNC:
             server = TCPServer(("localhost", 9999), MockTCPHandler)
             server.events = events
             server_thread = threading.Thread(target=server.handle_request_and_shutdown)
             server_thread.start()
             _c = self.simple_client(
-                "mongodb://localhost:9999", serverSelectionTimeoutMS=500, event_listeners=(listener,)
+                "mongodb://localhost:9999",
+                serverSelectionTimeoutMS=500,
+                event_listeners=(listener,),
             )
             server_thread.join()
             listener.wait_for_event(ServerHeartbeatStartedEvent, 1)
             listener.wait_for_event(ServerHeartbeatFailedEvent, 1)
 
-            self.assertEqual(
-                events,
-                [
-                    "serverHeartbeatStartedEvent",
-                    "client connected",
-                    "client hello received",
-                    "serverHeartbeatFailedEvent",
-                ],
-            )
-    else:
-        def test_heartbeat_start_ordering(self):
-            events = []
+        else:
 
-            def handle_client(reader: StreamReader, writer):
-                server.events.append("client connected")
-                print("clent connected")
+            def handle_client(reader: StreamReader, writer: StreamWriter):
+                events.append("client connected")
                 if (reader.read(1024)).strip():
-                    server.events.append("client hello received")
-                    print("client helllo recieved")
-            listener = HeartbeatEventsListListener(events)
+                    events.append("client hello received")
+                writer.close()
+                writer.wait_closed()
+
             server = asyncio.start_server(handle_client, "localhost", 9999)
-            with server:
-                server.events = events
-                _c = self.simple_client(
-                    "mongodb://localhost:9999", serverSelectionTimeoutMS=500, event_listeners=(listener,)
-                )
-                server.close()
-                server_task = asyncio.create_task(server.wait_closed())
-                server_task
-                listener.wait_for_event(ServerHeartbeatStartedEvent, 1)
-                listener.wait_for_event(ServerHeartbeatFailedEvent, 1)
-
-                self.assertEqual(
-                    events,
-                    [
-                        "serverHeartbeatStartedEvent",
-                        "client connected",
-                        "client hello received",
-                        "serverHeartbeatFailedEvent",
-                    ],
-                )
+            server.events = events
+            server.start_serving()
+            print(server.is_serving())
+            _c = self.simple_client(
+                "mongodb://localhost:9999",
+                serverSelectionTimeoutMS=500,
+                event_listeners=(listener,),
+            )
+            if _c._options.connect:
+                _c._connect()
+
+            listener.wait_for_event(ServerHeartbeatStartedEvent, 1)
+            listener.wait_for_event(ServerHeartbeatFailedEvent, 1)
+
+            server.close()
+            server.wait_closed()
+            _c.close()
+
+        self.assertEqual(
+            events,
+            [
+                "serverHeartbeatStartedEvent",
+                "client connected",
+                "client hello received",
+                "serverHeartbeatFailedEvent",
+            ],
+        )
 
 
 # Generate unified tests.