From 299f02d36c4ef1f7f44046c2493a6adb3b588452 Mon Sep 17 00:00:00 2001
From: Ben Johnson <ben@binarylogic.com>
Date: Thu, 11 Apr 2024 14:09:12 -0700
Subject: [PATCH] Handle CancelledError in an outer try block

Signed-off-by: Ben Johnson <ben@binarylogic.com>
---
 tests/test_trinnov_altitude.py       |  5 +++
 trinnov_altitude/trinnov_altitude.py | 66 +++++++++++++++-------------
 2 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/tests/test_trinnov_altitude.py b/tests/test_trinnov_altitude.py
index f38eff8..c36e943 100644
--- a/tests/test_trinnov_altitude.py
+++ b/tests/test_trinnov_altitude.py
@@ -100,6 +100,11 @@ async def test_start_listening_reconnects(mock_server):
 # --------------------------
 
 
+@pytest.mark.asyncio
+async def test_power_off(mock_server, connected_client):
+    await connected_client.power_off()
+
+
 @pytest.mark.asyncio
 async def test_volume_adjust(mock_server, connected_client):
     await connected_client.volume_adjust(2)
diff --git a/trinnov_altitude/trinnov_altitude.py b/trinnov_altitude/trinnov_altitude.py
index f2864e4..dfecdbe 100644
--- a/trinnov_altitude/trinnov_altitude.py
+++ b/trinnov_altitude/trinnov_altitude.py
@@ -135,7 +135,7 @@ def deregister_callback(self, callback: Callback):
     async def disconnect(self, timeout: int | float | None = USE_DEFAULT_TIMEOUT):
         """Closes the TCP connection to the processor"""
         if self._writer is None:
-            self.logger.warn("Not connected to Trinnov Altitude, can't disconnect")
+            self.logger.warning("Not connected to Trinnov Altitude, can't disconnect")
             return
 
         if timeout is self.USE_DEFAULT_TIMEOUT:
@@ -188,6 +188,9 @@ def start_listening(
 
     async def stop_listening(self, timeout: int | float | None = USE_DEFAULT_TIMEOUT):
         if self._response_handler_task:
+            if timeout is self.USE_DEFAULT_TIMEOUT:
+                timeout = self.timeout
+
             self._response_handler_task.cancel()
             await asyncio.wait_for(self._response_handler_task, timeout)
             self._response_handler_task = None
@@ -567,37 +570,40 @@ async def _listen(self, reconnect: bool, reconnect_backoff: float):
         This method will automatically reconnect when necessary if `reconnect`
         is set to `True`.
         """
-        failure_count = 0
-
-        while True:
-            try:
-                await self._read(timeout=None)
-            except (exceptions.NotConnectedError, EOFError, OSError) as e:
-                failure_count += 1
-
-                if reconnect:
-                    level = logging.DEBUG if failure_count <= 1 else logging.ERROR
-                    self.logger.log(
-                        level,
-                        f"Unable to read message from Trinnov Altitude, reconnecting...: {e}",
-                    )
-
-                    try:
-                        await self.reconnect(timeout=None)
-                    except Exception as err:  # pylint: disable=broad-except
-                        self.logger.exception(
-                            "Unhandled exception %s('%s')", type(err).__name__, err
+        try:
+            failure_count = 0
+
+            while True:
+                try:
+                    await self._read(timeout=None)
+                except (exceptions.NotConnectedError, EOFError, OSError) as e:
+                    failure_count += 1
+
+                    if reconnect:
+                        level = logging.DEBUG if failure_count <= 1 else logging.ERROR
+                        self.logger.log(
+                            level,
+                            f"Unable to read message from Trinnov Altitude, reconnecting...: {e}",
                         )
-                        await asyncio.sleep(reconnect_backoff)
+
+                        try:
+                            await self.reconnect(timeout=None)
+                        except (
+                            exceptions.ConnectionTimeoutError,
+                            exceptions.ConnectionFailedError,
+                        ) as e:
+                            self.logger.debug(
+                                f"Trinnov Altitude reconnect failed, trying again in {reconnect_backoff} seconds...: {e}"
+                            )
+                            await asyncio.sleep(reconnect_backoff)
+                    else:
+                        raise e
                 else:
-                    raise e
-            except asyncio.CancelledError:
-                self.logger.debug(
-                    "Trinnov Altitude listen task received cancel, shutting down..."
-                )
-                break
-            else:
-                failure_count = 0
+                    failure_count = 0
+        except asyncio.CancelledError:
+            self.logger.error(
+                "Trinnov Altitude listen task received cancel, shutting down..."
+            )
 
     def _process_message(self, raw_message: str):  # noqa: C901
         """Receive a single message off of the socket and process it."""