From 3f18cdc75ef73ef58cd9cb7bbc82effb1b2923ba Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 25 Oct 2020 07:27:51 -0500
Subject: [PATCH] Refactor august to use self.async_on_remove (#42333)

---
 homeassistant/components/august/entity.py     | 21 +++++++------------
 homeassistant/components/august/subscriber.py | 10 ++++++++-
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/homeassistant/components/august/entity.py b/homeassistant/components/august/entity.py
index 0d259778796..b6c677a63b6 100644
--- a/homeassistant/components/august/entity.py
+++ b/homeassistant/components/august/entity.py
@@ -46,18 +46,13 @@ class AugustEntityMixin(Entity):
 
     async def async_added_to_hass(self):
         """Subscribe to updates."""
-        self._data.async_subscribe_device_id(
-            self._device_id, self._update_from_data_and_write_state
+        self.async_on_remove(
+            self._data.async_subscribe_device_id(
+                self._device_id, self._update_from_data_and_write_state
+            )
         )
-        self._data.activity_stream.async_subscribe_device_id(
-            self._device_id, self._update_from_data_and_write_state
-        )
-
-    async def async_will_remove_from_hass(self):
-        """Undo subscription."""
-        self._data.async_unsubscribe_device_id(
-            self._device_id, self._update_from_data_and_write_state
-        )
-        self._data.activity_stream.async_unsubscribe_device_id(
-            self._device_id, self._update_from_data_and_write_state
+        self.async_on_remove(
+            self._data.activity_stream.async_subscribe_device_id(
+                self._device_id, self._update_from_data_and_write_state
+            )
         )
diff --git a/homeassistant/components/august/subscriber.py b/homeassistant/components/august/subscriber.py
index 81538fa011e..3a7edd8a342 100644
--- a/homeassistant/components/august/subscriber.py
+++ b/homeassistant/components/august/subscriber.py
@@ -18,13 +18,21 @@ class AugustSubscriberMixin:
 
     @callback
     def async_subscribe_device_id(self, device_id, update_callback):
-        """Add an callback subscriber."""
+        """Add an callback subscriber.
+
+        Returns a callable that can be used to unsubscribe.
+        """
         if not self._subscriptions:
             self._unsub_interval = async_track_time_interval(
                 self._hass, self._async_refresh, self._update_interval
             )
         self._subscriptions.setdefault(device_id, []).append(update_callback)
 
+        def _unsubscribe():
+            self.async_unsubscribe_device_id(device_id, update_callback)
+
+        return _unsubscribe
+
     @callback
     def async_unsubscribe_device_id(self, device_id, update_callback):
         """Remove a callback subscriber."""