Added timeout parameter to websocket.send.

This commit is contained in:
Fabio Manganiello 2023-02-15 22:23:15 +01:00
parent a7a107e5fb
commit 08553f84b9
Signed by untrusted user: blacklight
GPG key ID: D90FBA7F76362774

View file

@ -50,6 +50,7 @@ class WebsocketPlugin(AsyncRunnablePlugin):
ssl_cafile=None, ssl_cafile=None,
ssl_capath=None, ssl_capath=None,
wait_response=False, wait_response=False,
timeout=None,
): ):
""" """
Sends a message to a websocket. Sends a message to a websocket.
@ -66,6 +67,9 @@ class WebsocketPlugin(AsyncRunnablePlugin):
required by the SSL configuration (default: None) required by the SSL configuration (default: None)
:param wait_response: Set to True if you expect a response to the :param wait_response: Set to True if you expect a response to the
delivered message. delivered message.
:param timeout: If ``wait_response=True``, then ``timeout`` establishes
how long we should wait for a response before returning (default:
no timeout).
:return: The received response if ``wait_response`` is set to True, :return: The received response if ``wait_response`` is set to True,
otherwise nothing. otherwise nothing.
""" """
@ -87,10 +91,13 @@ class WebsocketPlugin(AsyncRunnablePlugin):
try: try:
await ws.send(str(msg)) await ws.send(str(msg))
except ConnectionClosed as err: except ConnectionClosed as err:
self.logger.warning('Error on websocket %s: %s', url, err) self.logger.warning(
'Connection error to websocket %s: %s', url, err
)
if wait_response: if wait_response:
messages = await self._recv(ws, num_messages=1) messages = await self._recv(ws, num_messages=1, timeout=timeout)
if messages: if messages:
return self._parse_msg(messages[0]) return self._parse_msg(messages[0])
@ -145,7 +152,7 @@ class WebsocketPlugin(AsyncRunnablePlugin):
return self.loop.call_soon_threadsafe(recv) return self.loop.call_soon_threadsafe(recv)
async def _recv(self, ws, timeout=0, num_messages=0): async def _recv(self, ws, timeout: Optional[float] = 0, num_messages=0):
messages = [] messages = []
time_start = time.time() time_start = time.time()
time_end = time_start + timeout if timeout else 0 time_end = time_start + timeout if timeout else 0