diff --git a/custom_components/husqvarna_automower/__init__.py b/custom_components/husqvarna_automower/__init__.py
index 1d18671a..a5e32df5 100644
--- a/custom_components/husqvarna_automower/__init__.py
+++ b/custom_components/husqvarna_automower/__init__.py
@@ -2,7 +2,6 @@
 import logging
 
 import aioautomower
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_PASSWORD, CONF_TOKEN, CONF_USERNAME
 from homeassistant.core import HomeAssistant
diff --git a/custom_components/husqvarna_automower/config_flow.py b/custom_components/husqvarna_automower/config_flow.py
index 82f386aa..51e6ec57 100644
--- a/custom_components/husqvarna_automower/config_flow.py
+++ b/custom_components/husqvarna_automower/config_flow.py
@@ -1,10 +1,10 @@
 """Config flow to add the integration via the UI."""
 import logging
 
-from aioautomower import GetAccessToken, GetMowerData, TokenError
 from aiohttp.client_exceptions import ClientConnectorError, ClientResponseError
 import voluptuous as vol
 
+from aioautomower import GetAccessToken, GetMowerData, TokenError
 from homeassistant import config_entries
 from homeassistant.const import (
     CONF_ACCESS_TOKEN,
diff --git a/custom_components/husqvarna_automower/entity.py b/custom_components/husqvarna_automower/entity.py
index f966eeaa..bf510f92 100644
--- a/custom_components/husqvarna_automower/entity.py
+++ b/custom_components/husqvarna_automower/entity.py
@@ -23,10 +23,6 @@ def __init__(self, session, idx) -> None:
         self.mower_name = mower_attributes["system"]["name"]
         self.model = mower_attributes["system"]["model"]
 
-        self.session.register_cb(
-            lambda _: self.async_write_ha_state(), schedule_immediately=True
-        )
-
         self._available = self.get_mower_attributes()["metadata"]["connected"]
 
         self._event = None
@@ -36,6 +32,18 @@ def get_mower_attributes(self) -> dict:
         """Get the mower attributes of the current mower."""
         return self.session.data["data"][self.idx]["attributes"]
 
+    async def async_added_to_hass(self) -> None:
+        """Call when entity about to be added to Home Assistant."""
+        await super().async_added_to_hass()
+        self.session.register_data_callback(
+            lambda _: self.async_write_ha_state(), schedule_immediately=True
+        )
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Call when entity is being removed from Home Assistant."""
+        await super().async_will_remove_from_hass()
+        self.session.unregister_data_callback(lambda _: self.async_write_ha_state())
+
     @property
     def device_info(self) -> DeviceInfo:
         """Defines the DeviceInfo for the mower."""
diff --git a/custom_components/husqvarna_automower/manifest.json b/custom_components/husqvarna_automower/manifest.json
index ea7a1aae..e7b5f9c3 100644
--- a/custom_components/husqvarna_automower/manifest.json
+++ b/custom_components/husqvarna_automower/manifest.json
@@ -11,7 +11,7 @@
     "@Thomas55555"
   ],
   "requirements": [
-    "aioautomower==2022.2.0"
+    "aioautomower==2022.2.1"
   ],
   "iot_class": "cloud_push",
   "version": "0.0.0"
diff --git a/custom_components/husqvarna_automower/system_health.py b/custom_components/husqvarna_automower/system_health.py
index 62f05bba..ff862ae9 100644
--- a/custom_components/husqvarna_automower/system_health.py
+++ b/custom_components/husqvarna_automower/system_health.py
@@ -1,6 +1,5 @@
 """Provide info to system health."""
 from aioautomower import TOKEN_URL
-
 from homeassistant.components.system_health import (
     SystemHealthRegistration,
     async_check_can_reach_url,