From eacf758900de9751bb6d9535167e093354ed9d52 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 13 Mar 2023 09:56:36 -1000 Subject: [PATCH] Allow any message from the remote to cancel the pong timer (#401) --- aioesphomeapi/connection.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py index 2aa8a43..f9c0c9a 100644 --- a/aioesphomeapi/connection.py +++ b/aioesphomeapi/connection.py @@ -51,7 +51,7 @@ _LOGGER = logging.getLogger(__name__) BUFFER_SIZE = 1024 * 1024 # Set buffer limit to 1MB -INTERNAL_MESSAGE_TYPES = {GetTimeRequest, PingRequest, PingResponse, DisconnectRequest} +INTERNAL_MESSAGE_TYPES = {GetTimeRequest, PingRequest, DisconnectRequest} PING_REQUEST_MESSAGE = PingRequest() PING_RESPONSE_MESSAGE = PingResponse() @@ -679,6 +679,11 @@ class APIConnection: _LOGGER.debug("%s: Got message of type %s: %s", self.log_name, msg_type, msg) + if self._pong_timer: + # Any valid message from the remote cancels the pong timer + # as we know the connection is still alive + self._async_cancel_pong_timer() + for handler in self._message_handlers.get(msg_type, [])[:]: handler(msg) @@ -687,11 +692,7 @@ class APIConnection: if msg_type not in INTERNAL_MESSAGE_TYPES: return - if msg_type is PingResponse: - # We got a pong so we know the ESP is alive, cancel the timer - # that will disconnect us - self._async_cancel_pong_timer() - elif msg_type is DisconnectRequest: + if msg_type is DisconnectRequest: self.send_message(DisconnectResponse()) self._connection_state = ConnectionState.CLOSED self._expected_disconnect = True