From 93cd596152932a6b098180e556e3bb5f2acc7373 Mon Sep 17 00:00:00 2001 From: Taddes Date: Fri, 8 Mar 2024 18:32:47 -0500 Subject: [PATCH] asyncio.timeout context managers --- tests/integration/async_push_test_client.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/tests/integration/async_push_test_client.py b/tests/integration/async_push_test_client.py index 5be7f5013..a8c9f7182 100644 --- a/tests/integration/async_push_test_client.py +++ b/tests/integration/async_push_test_client.py @@ -233,31 +233,27 @@ async def send_notification( else: return resp - def get_notification(self, timeout=1): + async def get_notification(self, timeout=1): """Get notification.""" if not self.ws: raise WebSocketException("WebSocket client not available as expected") - orig_timeout = self.ws.time - self.ws.settimeout(timeout) try: - d = self.ws.recv() - log.debug(f"Recv: {d}") + async with asyncio.timeout(timeout): + d = await self.ws.recv() + log.debug(f"Recv: {d!r}") return json.loads(d) except Exception: return None - finally: - self.ws.settimeout(orig_timeout) - def get_broadcast(self, timeout=1): # pragma: no cover + async def get_broadcast(self, timeout=1): # pragma: no cover """Get broadcast.""" if not self.ws: raise WebSocketException("WebSocket client not available as expected") - orig_timeout = self.ws.gettimeout() - self.ws.settimeout(timeout) try: - d = self.ws.recv() + async with asyncio.timeout(timeout): + d = await self.ws.recv() log.debug(f"Recv: {d}") result = json.loads(d) # ASK JR @@ -266,8 +262,6 @@ def get_broadcast(self, timeout=1): # pragma: no cover except WebSocketException as ex: # pragma: no cover log.error(f"Error: {ex}") return None - finally: - self.ws.settimeout(orig_timeout) async def ping(self): """Test ping."""