cleanup registered callbacks before removing velux config entry (home-assistant#156525)

Author: wollew

homeassistant/components/velux/entity.py

@@ -1,8 +1,9 @@
 """Support for VELUX KLF 200 devices."""
 
+from collections.abc import Awaitable, Callable
+
 from pyvlx import Node
 
-from homeassistant.core import callback
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
 
@@ -14,6 +15,7 @@ class VeluxEntity(Entity):
 
     _attr_should_poll = False
     _attr_has_entity_name = True
+    update_callback: Callable[["Node"], Awaitable[None]] | None = None
 
     def __init__(self, node: Node, config_entry_id: str) -> None:
         """Initialize the Velux device."""
@@ -24,6 +26,7 @@ def __init__(self, node: Node, config_entry_id: str) -> None:
             else f"{config_entry_id}_{node.node_id}"
         )
         self._attr_unique_id = unique_id
+        self.unsubscribe = None
 
         self._attr_device_info = DeviceInfo(
             identifiers={
@@ -37,16 +40,18 @@ def __init__(self, node: Node, config_entry_id: str) -> None:
             via_device=(DOMAIN, f"gateway_{config_entry_id}"),
         )
 
-    @callback
-    def async_register_callbacks(self):
-        """Register callbacks to update hass after device was changed."""
+    async def after_update_callback(self, node) -> None:
+        """Call after device was updated."""
+        self.async_write_ha_state()
 
-        async def after_update_callback(device):
-            """Call after device was updated."""
-            self.async_write_ha_state()
+    async def async_added_to_hass(self) -> None:
+        """Register callback and store reference for cleanup."""
 
-        self.node.register_device_updated_cb(after_update_callback)
+        self.update_callback = self.after_update_callback
+        self.node.register_device_updated_cb(self.update_callback)
 
-    async def async_added_to_hass(self) -> None:
-        """Store register state change callback."""
-        self.async_register_callbacks()
+    async def async_will_remove_from_hass(self) -> None:
+        """Clean up registered callbacks."""
+        if self.update_callback:
+            self.node.unregister_device_updated_cb(self.update_callback)
+            self.update_callback = None

homeassistant/components/velux/quality_scale.yaml

@@ -15,9 +15,7 @@ rules:
   docs-high-level-description: done
   docs-installation-instructions: done
   docs-removal-instructions: done
-  entity-event-setup:
-    status: todo
-    comment: subscribe is ok, unsubscribe needs to be added
+  entity-event-setup: done
   entity-unique-id: done
   has-entity-name:
     status: todo

tests/components/velux/test_light.py

@@ -8,6 +8,8 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
+from tests.common import MockConfigEntry
+
 
 @pytest.fixture
 def platform() -> Platform:
@@ -41,3 +43,30 @@ async def test_light_setup(
     # Verify device has correct identifiers + name
     assert ("velux", mock_light.serial_number) in device_entry.identifiers
     assert device_entry.name == mock_light.name
+
+
+# This test is not light specific, it just uses the light platform to test the base entity class.
+@pytest.mark.usefixtures("setup_integration")
+async def test_entity_callbacks(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_light: AsyncMock,
+) -> None:
+    """Ensure the entity unregisters its device-updated callback when unloaded."""
+    # Entity is created by setup_integration; callback should be registered
+    test_entity_id = f"light.{mock_light.name.lower().replace(' ', '_')}"
+    state = hass.states.get(test_entity_id)
+    assert state is not None
+
+    # Callback is registered exactly once with a callable
+    assert mock_light.register_device_updated_cb.call_count == 1
+    cb = mock_light.register_device_updated_cb.call_args[0][0]
+    assert callable(cb)
+
+    # Unload the config entry to trigger async_will_remove_from_hass
+    assert await hass.config_entries.async_unload(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Callback must be unregistered with the same callable
+    assert mock_light.unregister_device_updated_cb.call_count == 1
+    assert mock_light.unregister_device_updated_cb.call_args[0][0] is cb