add on_disconnected_event

Author: MoChilia

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_auto_connect.py

@@ -25,7 +25,7 @@ def test_auto_connect(self, webpubsubclient_endpoint):
             reconnect_retry_backoff_factor=0.1,
         )
         name = "test_auto_connect"
-        connected_event, message_event = self.setup_events(client)
+        connected_event, _, message_event = self.setup_events(client)
         with client:
             assert connected_event.wait(timeout=30), "Timed out waiting for initial connection"
             conn_id0 = client._connection_id

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_auto_connect_async.py

@@ -26,7 +26,7 @@ async def test_auto_connect_async(self, webpubsubclient_endpoint):
             reconnect_retry_backoff_factor=0.1,
         )
         name = "test_auto_connect_async"
-        connected_event, message_event = await self.setup_events(client)
+        connected_event, _, message_event = await self.setup_events(client)
         async with client:
             await asyncio.wait_for(connected_event.wait(), timeout=30)
             conn_id0 = client._connection_id

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_no_recovery_no_connect.py

@@ -34,14 +34,10 @@ def test_disable_recovery_and_autoconnect(self, webpubsubclient_endpoint):
         name = "test_disable_recovery_and_autoconnect"
         with client:
             group_name = name
-            client.subscribe("group-message", on_group_message)
+            _, disconnected_event, _ = self.setup_events(client)
             client.join_group(group_name)
             client._ws.sock.close(1001)  # close connection
-            # wait for client to detect disconnection so send raises SendMessageError
-            for _ in range(30):
-                if not client.is_connected():
-                    break
-                time.sleep(1)
+            assert disconnected_event.wait(timeout=30), "Timed out waiting for disconnection"
             with pytest.raises(SendMessageError):
                 client.send_to_group(group_name, name, "text")
             time.sleep(3)  # wait to confirm message was NOT received
@@ -63,12 +59,10 @@ def test_disable_recovery_and_autoconnect_send_concurrently(
 
         with client:
             group_name = "test_disable_recovery_and_autoconnect_send_concurrently"
+            _, disconnected_event, _ = self.setup_events(client)
             client.join_group(group_name)
             client._ws.sock.close(1001)  # close connection
-            for _ in range(30):
-                if not client.is_connected():
-                    break
-                time.sleep(1)
+            assert disconnected_event.wait(timeout=30), "Timed out waiting for disconnection"
             assert not client.is_connected()
 
             def send(idx):

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_no_recovery_no_connect_async.py

@@ -29,14 +29,10 @@ async def test_disable_recovery_and_autoconnect_async(self, webpubsubclient_endp
         name = "test_disable_recovery_and_autoconnect_async"
         async with client:
             group_name = name
-            await client.subscribe("group-message", on_group_message)
+            _, disconnected_event, _ = await self.setup_events(client)
             await client.join_group(group_name)
             await client._ws.session.close()  # close connection
-            # wait for client to detect disconnection so send raises SendMessageError
-            for _ in range(30):
-                if not client.is_connected():
-                    break
-                await asyncio.sleep(1)
+            await asyncio.wait_for(disconnected_event.wait(), timeout=30)
             with pytest.raises(SendMessageError):
                 await client.send_to_group(group_name, name, "text")
             await asyncio.sleep(3)  # wait to confirm message was NOT received
@@ -58,13 +54,10 @@ async def test_disable_recovery_and_autoconnect_send_concurrently_async(
 
         async with client:
             group_name = "test_disable_recovery_and_autoconnect_send_concurrently_async"
+            _, disconnected_event, _ = await self.setup_events(client)
             await client.join_group(group_name)
             await client._ws.session.close()  # close connection
-            # wait for client to detect disconnection
-            for _ in range(30):
-                if not client.is_connected():
-                    break
-                await asyncio.sleep(1)
+            await asyncio.wait_for(disconnected_event.wait(), timeout=30)
             assert not client.is_connected()
 
             tasks = [client.send_to_group(group_name, "hello", "text") for _ in range(10)]

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_recovery.py

@@ -17,7 +17,7 @@ class TestWebpubsubClientRecovery(WebpubsubClientTest):
     def test_recovery(self, webpubsubclient_endpoint):
         client = self.create_client(endpoint=webpubsubclient_endpoint, message_retry_total=10)
         name = "test_recovery"
-        connected_event, message_event = self.setup_events(client)
+        connected_event, _, message_event = self.setup_events(client)
         with client:
             assert connected_event.wait(timeout=30), "Timed out waiting for connection"
             conn_id0 = client._connection_id

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_recovery_async.py

@@ -19,7 +19,7 @@ class TestWebpubsubClientRecoveryAsync(WebpubsubClientTestAsync):
     async def test_recovery_async(self, webpubsubclient_endpoint):
         client = await self.create_client(endpoint=webpubsubclient_endpoint, message_retry_total=10)
         name = "test_recovery_async"
-        connected_event, message_event = await self.setup_events(client)
+        connected_event, _, message_event = await self.setup_events(client)
         async with client:
             await asyncio.wait_for(connected_event.wait(), timeout=30)
             conn_id0 = client._connection_id

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_send_concurrently.py

@@ -4,7 +4,6 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # -------------------------------------------------------------------------
-import time
 import pytest
 from devtools_testutils import recorded_by_proxy
 from testcase import WebpubsubClientTest, WebpubsubClientPowerShellPreparer, SafeThread
@@ -16,14 +15,11 @@ class TestWebpubsubClientSendConcurrently(WebpubsubClientTest):
     @recorded_by_proxy
     def test_send_concurrently(self, webpubsubclient_endpoint):
         client = self.create_client(endpoint=webpubsubclient_endpoint)
+        connected_event, _, _ = self.setup_events(client)
         with client:
+            assert connected_event.wait(timeout=30), "Timed out waiting for connection"
             group_name = "test_send_concurrently"
             client.join_group(group_name)
-            # wait for connection to stabilize before concurrent sends
-            for _ in range(30):
-                if client.is_connected():
-                    break
-                time.sleep(1)
 
             def send(idx):
                 client.send_to_group(group_name, f"hello_{idx}", "text")

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_send_concurrently_async.py

@@ -17,11 +17,9 @@ class TestWebpubsubClientSendConcurrentlyAsync(WebpubsubClientTestAsync):
     @recorded_by_proxy_async
     async def test_send_concurrently_async(self, webpubsubclient_endpoint):
         client = await self.create_client(endpoint=webpubsubclient_endpoint)
+        connected_event, _, _ = await self.setup_events(client)
         async with client:
+            await asyncio.wait_for(connected_event.wait(), timeout=30)
             group_name = "test_send_concurrently_async"
             await client.join_group(group_name)
-            for _ in range(30):
-                if client.is_connected():
-                    break
-                await asyncio.sleep(1)
             await asyncio.gather(*[client.send_to_group(group_name, f"hello_{idx}", "text") for idx in range(100)])

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_smoke.py

@@ -29,38 +29,44 @@ class TestWebpubsubClientSmoke(WebpubsubClientTest):
     def test_call_back_deadlock(self, webpubsubclient_endpoint):
         client = self.create_client(endpoint=webpubsubclient_endpoint)
         group_name = "test_call_back_deadlock"
+        callback_completed = threading.Event()
+        callback_count = 0
+        callback_count_lock = threading.Lock()
 
         def on_group_message(msg: OnGroupDataMessageArgs):
+            nonlocal callback_count
             client.send_to_group(group_name, msg.data, "text", no_echo=True)
+            with callback_count_lock:
+                callback_count += 1
+                if callback_count >= 3:
+                    callback_completed.set()
 
         with client:
             client.join_group(group_name)
             client.subscribe("group-message", on_group_message)
             client.send_to_group(group_name, "hello test_call_back_deadlock1", "text")
             client.send_to_group(group_name, "hello test_call_back_deadlock2", "text")
             client.send_to_group(group_name, "hello test_call_back_deadlock3", "text")
-            # sleep to make sure the callback has enough time to execute before close
-            time.sleep(1)
+            assert callback_completed.wait(timeout=30), "Timed out waiting for callbacks to finish"
 
     @WebpubsubClientPowerShellPreparer()
     @recorded_by_proxy
     def test_context_manager(self, webpubsubclient_endpoint):
         client = self.create_client(endpoint=webpubsubclient_endpoint)
+        _, _, message_event = self.setup_events(client)
         with client:
             group_name = "test_context_manager"
             client.join_group(group_name)
             client.send_to_group(group_name, "test_context_manager", "text")
-            for _ in range(30):
-                if client._sequence_id.sequence_id > 0:
-                    break
-                time.sleep(1)
+            assert message_event.wait(timeout=30), "Timed out waiting for context manager message"
             assert client._sequence_id.sequence_id > 0
 
     # test on_stop
     @WebpubsubClientPowerShellPreparer()
     @recorded_by_proxy
     def test_on_stop(self, webpubsubclient_endpoint):
         client = self.create_client(endpoint=webpubsubclient_endpoint)
+        connected_event, disconnected_event, _ = self.setup_events(client)
         reopen_error = None
         reopen_complete = threading.Event()
 
@@ -77,10 +83,8 @@ def on_stop():
                 else:
                     raise RuntimeError("Failed to reopen client in stopped callback")
 
-                for _ in range(30):
-                    if client.is_connected():
-                        break
-                    time.sleep(1)
+                if not connected_event.wait(timeout=30):
+                    raise RuntimeError("Timed out waiting for client to reconnect in stopped callback")
                 assert client.is_connected()
             except Exception as e:
                 reopen_error = e
@@ -90,11 +94,9 @@ def on_stop():
         with client:
             # open client again after close
             client.subscribe("stopped", on_stop)
-            for _ in range(30):
-                if client.is_connected():
-                    break
-                time.sleep(1)
+            assert connected_event.wait(timeout=30), "Timed out waiting for initial connection"
             assert client.is_connected()
+            connected_event.clear()
             client.close()
             # wait for on_stop callback to finish reopening
             if not reopen_complete.wait(timeout=60):
@@ -103,12 +105,9 @@ def on_stop():
 
             # remove stopped event and close again
             client.unsubscribe("stopped", on_stop)
+            disconnected_event.clear()
             client.close()
-            # wait for disconnect to finalize
-            for _ in range(30):
-                if not client.is_connected():
-                    break
-                time.sleep(1)
+            assert disconnected_event.wait(timeout=30), "Timed out waiting for client to disconnect"
             assert not client.is_connected()
 
     @WebpubsubClientPowerShellPreparer()
@@ -148,7 +147,7 @@ def _test(enable_auto_rejoin, test_group_name, assert_func):
                 auto_rejoin_groups=enable_auto_rejoin,
                 message_retry_total=10,
             )
-            connected_event, message_event = self.setup_events(client)
+            connected_event, _, message_event = self.setup_events(client)
             with client:
                 client.join_group(test_group_name)
 

sdk/webpubsub/azure-messaging-webpubsubclient/tests/test_smoke_async.py

@@ -28,38 +28,42 @@ class TestWebpubsubClientSmokeAsync(WebpubsubClientTestAsync):
     async def test_call_back_deadlock_async(self, webpubsubclient_endpoint):
         client = await self.create_client(endpoint=webpubsubclient_endpoint)
         group_name = "test_call_back_deadlock_async"
+        callback_completed = asyncio.Event()
+        callback_count = 0
 
         async def on_group_message(msg: OnGroupDataMessageArgs):
+            nonlocal callback_count
             await client.send_to_group(group_name, msg.data, "text", no_echo=True)
+            callback_count += 1
+            if callback_count >= 3:
+                callback_completed.set()
 
         async with client:
             await client.join_group(group_name)
             await client.subscribe("group-message", on_group_message)
             await client.send_to_group(group_name, "hello test_call_back_deadlock1", "text")
             await client.send_to_group(group_name, "hello test_call_back_deadlock2", "text")
             await client.send_to_group(group_name, "hello test_call_back_deadlock3", "text")
-            # sleep to make sure the callback has enough time to execute before close
-            await asyncio.sleep(1)
+            await asyncio.wait_for(callback_completed.wait(), timeout=30)
 
     @WebpubsubClientPowerShellPreparer()
     @recorded_by_proxy_async
     async def test_context_manager_async(self, webpubsubclient_endpoint):
         client = await self.create_client(endpoint=webpubsubclient_endpoint)
+        _, _, message_event = await self.setup_events(client)
         async with client:
             group_name = "test_context_manager_async"
             await client.join_group(group_name)
             await client.send_to_group(group_name, "test_context_manager", "text")
-            for _ in range(30):
-                if client._sequence_id.sequence_id > 0:
-                    break
-                await asyncio.sleep(1)
+            await asyncio.wait_for(message_event.wait(), timeout=30)
             assert client._sequence_id.sequence_id > 0
 
     # test on_stop
     @WebpubsubClientPowerShellPreparer()
     @recorded_by_proxy_async
     async def test_on_stop_async(self, webpubsubclient_endpoint):
         client = await self.create_client(endpoint=webpubsubclient_endpoint)
+        connected_event, disconnected_event, _ = await self.setup_events(client)
         reopen_error = None
         reopen_complete = asyncio.Event()
 
@@ -76,10 +80,7 @@ async def on_stop():
                 else:
                     raise RuntimeError("Failed to reopen client in stopped callback")
 
-                for _ in range(30):
-                    if client.is_connected():
-                        break
-                    await asyncio.sleep(1)
+                await asyncio.wait_for(connected_event.wait(), timeout=30)
                 assert client.is_connected()
             except Exception as e:
                 reopen_error = e
@@ -89,24 +90,19 @@ async def on_stop():
         async with client:
             # open client again after close
             await client.subscribe("stopped", on_stop)
-            for _ in range(30):
-                if client.is_connected():
-                    break
-                await asyncio.sleep(1)
+            await asyncio.wait_for(connected_event.wait(), timeout=30)
             assert client.is_connected()
+            connected_event.clear()
             await client.close()
             # wait for on_stop callback to finish reopening
             await asyncio.wait_for(reopen_complete.wait(), timeout=60)
             assert reopen_error is None, f"on_stop callback failed: {reopen_error}"
 
             # remove stopped event and close again
             await client.unsubscribe("stopped", on_stop)
+            disconnected_event.clear()
             await client.close()
-            # wait for disconnect to finalize
-            for _ in range(30):
-                if not client.is_connected():
-                    break
-                await asyncio.sleep(1)
+            await asyncio.wait_for(disconnected_event.wait(), timeout=30)
             assert not client.is_connected()
 
     @WebpubsubClientPowerShellPreparer()
@@ -145,7 +141,7 @@ async def _test(enable_auto_rejoin, test_group_name, assert_func):
                 endpoint=webpubsubclient_endpoint,
                 auto_rejoin_groups=enable_auto_rejoin,
             )
-            connected_event, message_event = await self.setup_events(client)
+            connected_event, _, message_event = await self.setup_events(client)
             async with client:
                 await client.join_group(test_group_name)