Fix entity removal leaks in wilight (#183867)

homeassistant/components/wilight/entity.py

@@ -1,11 +1,16 @@
 """The WiLight integration."""
 
+from functools import partial
 from typing import Any, override
 
 from pywilight.wilight_device import PyWiLightDevice
 
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
 from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN
@@ -27,6 +32,8 @@ class WiLightDevice(Entity):
         self._client = api_device.client
         self._index = index
         self._status: dict[str, Any] = {}
+        self._status_signal = f"wilight_status_{self._device_id}_{index}"
+        self._status_callback_registered = False
 
         self._attr_unique_id = f"{self._device_id}_{index}"
         self._attr_device_info = DeviceInfo(
@@ -56,5 +63,17 @@ class WiLightDevice(Entity):
     @override
     async def async_added_to_hass(self) -> None:
         """Register update callback."""
-        self._client.register_status_callback(self.handle_event_callback, self._index)
+        # pywilight can't unregister status callbacks, so register a forwarder
+        # once and (un)subscribe the entity through the dispatcher.
+        if not self._status_callback_registered:
+            self._client.register_status_callback(
+                partial(async_dispatcher_send, self.hass, self._status_signal),
+                self._index,
+            )
+            self._status_callback_registered = True
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass, self._status_signal, self.handle_event_callback
+            )
+        )
         await self._client.status(self._index)

tests/components/wilight/test_switch.py

@@ -4,6 +4,7 @@ from unittest.mock import patch
 
 import pytest
 import pywilight
+from pywilight.wilight_device import PyWiLightDevice
 
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.components.wilight.const import DOMAIN
@@ -40,6 +41,9 @@ from . import (
     setup_integration,
 )
 
+WATERING_ENTITY_ID = "switch.wl000000000099_1_watering"
+WATERING_INDEX = "0"
+
 
 @pytest.fixture(name="dummy_device_from_host_switch")
 def mock_dummy_device_from_host_switch():
@@ -261,3 +265,58 @@ async def test_switch_services(
         )
 
     assert str(exc_info.value) == "Entity is not a WiLight valve switch"
+
+
+async def test_status_callback_after_entity_removed(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    dummy_device_from_host_switch: PyWiLightDevice,
+) -> None:
+    """Test a removed entity no longer handles status callbacks."""
+    await setup_integration(hass)
+    client = dummy_device_from_host_switch.client
+
+    entity_registry.async_remove(WATERING_ENTITY_ID)
+    await hass.async_block_till_done()
+    assert hass.states.get(WATERING_ENTITY_ID) is None
+
+    with patch(
+        "homeassistant.components.wilight.entity.WiLightDevice.async_write_ha_state"
+    ) as mock_write_ha_state:
+        # Triggers the status callback registered with pywilight.
+        await client.turn_on(WATERING_INDEX)
+        await hass.async_block_till_done()
+
+    mock_write_ha_state.assert_not_called()
+    assert hass.states.get(WATERING_ENTITY_ID) is None
+
+
+async def test_status_callback_after_entity_readded(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    dummy_device_from_host_switch: PyWiLightDevice,
+) -> None:
+    """Test re-adding an entity does not register another status callback."""
+    await setup_integration(hass)
+    client = dummy_device_from_host_switch.client
+
+    with patch.object(
+        client, "register_status_callback", wraps=client.register_status_callback
+    ) as mock_register:
+        # Changing the entity_id removes and re-adds the same entity object.
+        entity_registry.async_update_entity(
+            WATERING_ENTITY_ID, new_entity_id="switch.renamed_watering"
+        )
+        await hass.async_block_till_done()
+
+    mock_register.assert_not_called()
+    assert hass.states.get(WATERING_ENTITY_ID) is None
+    state = hass.states.get("switch.renamed_watering")
+    assert state
+    assert state.state == STATE_OFF
+
+    await client.turn_on(WATERING_INDEX)
+    await hass.async_block_till_done()
+    state = hass.states.get("switch.renamed_watering")
+    assert state
+    assert state.state == STATE_ON