From b77fc2e8cbc4a89f7a46f401ea1a7fd5a60be876 Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Sun, 19 Dec 2021 15:05:13 -0700
Subject: [PATCH] Ensure existing SimpliSafe websocket tasks are cancelled
 appropriately (#62347)

---
 .../components/simplisafe/__init__.py         | 33 +++++++++++--------
 .../components/simplisafe/manifest.json       |  2 +-
 requirements_all.txt                          |  2 +-
 requirements_test_all.txt                     |  2 +-
 4 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/homeassistant/components/simplisafe/__init__.py b/homeassistant/components/simplisafe/__init__.py
index 537e199fc04..025c3a3ae49 100644
--- a/homeassistant/components/simplisafe/__init__.py
+++ b/homeassistant/components/simplisafe/__init__.py
@@ -520,14 +520,14 @@ class SimpliSafe:
         self._system_notifications[system.system_id] = latest_notifications
 
     async def _async_start_websocket_loop(self) -> None:
-        """Define a callback for connecting to the websocket."""
+        """Start a websocket reconnection loop."""
         if TYPE_CHECKING:
             assert self._api.websocket
 
         should_reconnect = True
 
         try:
-            await self._api.websocket.async_reconnect()
+            await self._api.websocket.async_connect()
             await self._api.websocket.async_listen()
         except asyncio.CancelledError:
             LOGGER.debug("Request to cancel websocket loop received")
@@ -539,10 +539,25 @@ class SimpliSafe:
 
         if should_reconnect:
             LOGGER.info("Disconnected from websocket; reconnecting")
+            await self._async_cancel_websocket_loop()
             self._websocket_reconnect_task = self._hass.async_create_task(
                 self._async_start_websocket_loop()
             )
 
+    async def _async_cancel_websocket_loop(self) -> None:
+        """Stop any existing websocket reconnection loop."""
+        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
+
+            if TYPE_CHECKING:
+                assert self._api.websocket
+            await self._api.websocket.async_disconnect()
+
     @callback
     def _async_websocket_on_event(self, event: WebsocketEvent) -> None:
         """Define a callback for receiving a websocket event."""
@@ -591,15 +606,7 @@ class SimpliSafe:
             if TYPE_CHECKING:
                 assert self._api.websocket
 
-            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()
+            await self._async_cancel_websocket_loop()
 
         self.entry.async_on_unload(
             self._hass.bus.async_listen_once(
@@ -641,8 +648,7 @@ class SimpliSafe:
                 data={**self.entry.data, CONF_TOKEN: token},
             )
 
-        @callback
-        def async_handle_refresh_token(token: str) -> None:
+        async def async_handle_refresh_token(token: str) -> None:
             """Handle a new refresh token."""
             async_save_refresh_token(token)
 
@@ -650,6 +656,7 @@ class SimpliSafe:
                 assert self._api.websocket
 
             # Open a new websocket connection with the fresh token:
+            await self._async_cancel_websocket_loop()
             self._websocket_reconnect_task = self._hass.async_create_task(
                 self._async_start_websocket_loop()
             )
diff --git a/homeassistant/components/simplisafe/manifest.json b/homeassistant/components/simplisafe/manifest.json
index 0b6cb385be6..8e494af013a 100644
--- a/homeassistant/components/simplisafe/manifest.json
+++ b/homeassistant/components/simplisafe/manifest.json
@@ -3,7 +3,7 @@
   "name": "SimpliSafe",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/simplisafe",
-  "requirements": ["simplisafe-python==2021.12.1"],
+  "requirements": ["simplisafe-python==2021.12.2"],
   "codeowners": ["@bachya"],
   "iot_class": "cloud_polling",
   "dhcp": [
diff --git a/requirements_all.txt b/requirements_all.txt
index 4ce995f10b9..3cde5400641 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -2158,7 +2158,7 @@ simplehound==0.3
 simplepush==1.1.4
 
 # homeassistant.components.simplisafe
-simplisafe-python==2021.12.1
+simplisafe-python==2021.12.2
 
 # homeassistant.components.sisyphus
 sisyphus-control==3.1.2
diff --git a/requirements_test_all.txt b/requirements_test_all.txt
index af3182af5c6..cfb9f5f8c76 100644
--- a/requirements_test_all.txt
+++ b/requirements_test_all.txt
@@ -1285,7 +1285,7 @@ sharkiqpy==0.1.8
 simplehound==0.3
 
 # homeassistant.components.simplisafe
-simplisafe-python==2021.12.1
+simplisafe-python==2021.12.2
 
 # homeassistant.components.slack
 slackclient==2.5.0