From 8eb33ede432a4defc832db9ec7f95e279b860ee9 Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Sun, 19 Dec 2021 13:52:21 -0700
Subject: [PATCH] Fix bug in which SimpliSafe websocket won't reconnect on
 error (#62241)

---
 .../components/simplisafe/__init__.py         | 48 +++++++++++++++----
 1 file changed, 38 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/simplisafe/__init__.py b/homeassistant/components/simplisafe/__init__.py
index e80716de9f8..537e199fc04 100644
--- a/homeassistant/components/simplisafe/__init__.py
+++ b/homeassistant/components/simplisafe/__init__.py
@@ -12,6 +12,7 @@ from simplipy.errors import (
     EndpointUnavailableError,
     InvalidCredentialsError,
     SimplipyError,
+    WebsocketError,
 )
 from simplipy.system import SystemNotification
 from simplipy.system.v3 import (
@@ -473,6 +474,7 @@ class SimpliSafe:
         self._api = api
         self._hass = hass
         self._system_notifications: dict[int, set[SystemNotification]] = {}
+        self._websocket_reconnect_task: asyncio.Task | None = None
         self.entry = entry
         self.initial_event_to_use: dict[int, dict[str, Any]] = {}
         self.systems: dict[int, SystemType] = {}
@@ -517,11 +519,29 @@ class SimpliSafe:
 
         self._system_notifications[system.system_id] = latest_notifications
 
-    async def _async_websocket_on_connect(self) -> None:
+    async def _async_start_websocket_loop(self) -> None:
         """Define a callback for connecting to the websocket."""
         if TYPE_CHECKING:
             assert self._api.websocket
-        await self._api.websocket.async_listen()
+
+        should_reconnect = True
+
+        try:
+            await self._api.websocket.async_reconnect()
+            await self._api.websocket.async_listen()
+        except asyncio.CancelledError:
+            LOGGER.debug("Request to cancel websocket loop received")
+            raise
+        except WebsocketError as err:
+            LOGGER.error("Failed to connect to websocket: %s", err)
+        except Exception as err:  # pylint: disable=broad-except
+            LOGGER.error("Unknown exception while connecting to websocket: %s", err)
+
+        if should_reconnect:
+            LOGGER.info("Disconnected from websocket; reconnecting")
+            self._websocket_reconnect_task = self._hass.async_create_task(
+                self._async_start_websocket_loop()
+            )
 
     @callback
     def _async_websocket_on_event(self, event: WebsocketEvent) -> None:
@@ -561,17 +581,25 @@ class SimpliSafe:
             assert self._api.refresh_token
             assert self._api.websocket
 
-        self._api.websocket.add_connect_callback(self._async_websocket_on_connect)
         self._api.websocket.add_event_callback(self._async_websocket_on_event)
-        asyncio.create_task(self._api.websocket.async_connect())
+        self._websocket_reconnect_task = asyncio.create_task(
+            self._async_start_websocket_loop()
+        )
 
         async def async_websocket_disconnect_listener(_: Event) -> None:
             """Define an event handler to disconnect from the websocket."""
             if TYPE_CHECKING:
                 assert self._api.websocket
 
-            if self._api.websocket.connected:
-                await self._api.websocket.async_disconnect()
+            if self._websocket_reconnect_task:
+                self._websocket_reconnect_task.cancel()
+                try:
+                    await self._websocket_reconnect_task
+                except asyncio.CancelledError:
+                    LOGGER.debug("Websocket reconnection task successfully canceled")
+                    self._websocket_reconnect_task = None
+
+            await self._api.websocket.async_disconnect()
 
         self.entry.async_on_unload(
             self._hass.bus.async_listen_once(
@@ -621,10 +649,10 @@ class SimpliSafe:
             if TYPE_CHECKING:
                 assert self._api.websocket
 
-            if self._api.websocket.connected:
-                # If a websocket connection is open, reconnect it to use the
-                # new access token:
-                asyncio.create_task(self._api.websocket.async_reconnect())
+            # Open a new websocket connection with the fresh token:
+            self._websocket_reconnect_task = self._hass.async_create_task(
+                self._async_start_websocket_loop()
+            )
 
         self.entry.async_on_unload(
             self._api.add_refresh_token_callback(async_handle_refresh_token)