diff --git a/homeassistant/components/knx/knx_entity.py b/homeassistant/components/knx/knx_entity.py index fff7f9b9f4f..9545510e635 100644 --- a/homeassistant/components/knx/knx_entity.py +++ b/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) diff --git a/homeassistant/components/knx/sensor.py b/homeassistant/components/knx/sensor.py index 4400c304193..dbfe8e9bd5e 100644 --- a/homeassistant/components/knx/sensor.py +++ b/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, + ) ) diff --git a/tests/components/knx/test_interface_device.py b/tests/components/knx/test_interface_device.py index 9fb21b9f9b4..12ae0ac7d0e 100644 --- a/tests/components/knx/test_interface_device.py +++ b/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,