Use `async_on_remove` for KNX entities removal (#95658)

* Use `async_on_remove` for KNX entities removal

* review

homeassistant/components/knx/knx_entity.py

@@ -42,8 +42,5 @@ class KnxEntity(Entity):
     async def async_added_to_hass(self) -> None:
         """Store register state change callback."""
         self._device.register_device_updated_cb(self.after_update_callback)
-
-    async def async_will_remove_from_hass(self) -> None:
-        """Disconnect device object when removed."""
-        # will also remove all callbacks
-        self._device.shutdown()
+        # will remove all callbacks and xknx tasks
+        self.async_on_remove(self._device.shutdown)

homeassistant/components/knx/sensor.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta
+from functools import partial
 from typing import Any
 
 from xknx import XKNX
@@ -221,9 +222,9 @@ class KNXSystemSensor(SensorEntity):
         self.knx.xknx.connection_manager.register_connection_state_changed_cb(
             self.after_update_callback
         )
-
-    async def async_will_remove_from_hass(self) -> None:
-        """Disconnect device object when removed."""
-        self.knx.xknx.connection_manager.unregister_connection_state_changed_cb(
-            self.after_update_callback
+        self.async_on_remove(
+            partial(
+                self.knx.xknx.connection_manager.unregister_connection_state_changed_cb,
+                self.after_update_callback,
+            )
         )

tests/components/knx/test_interface_device.py

@@ -99,11 +99,11 @@ async def test_removed_entity(
     hass: HomeAssistant, knx: KNXTestKit, entity_registry: er.EntityRegistry
 ) -> None:
     """Test unregister callback when entity is removed."""
-    await knx.setup_integration({})
-
-    with patch.object(
-        knx.xknx.connection_manager, "unregister_connection_state_changed_cb"
+    with patch(
+        "xknx.core.connection_manager.ConnectionManager.unregister_connection_state_changed_cb"
     ) as unregister_mock:
+        await knx.setup_integration({})
+
         entity_registry.async_update_entity(
             "sensor.knx_interface_connection_established",
             disabled_by=er.RegistryEntryDisabler.USER,