diff --git a/homeassistant/components/hue/hue_event.py b/homeassistant/components/hue/hue_event.py index ed1bc1c8f7d..a588f68ea4e 100644 --- a/homeassistant/components/hue/hue_event.py +++ b/homeassistant/components/hue/hue_event.py @@ -35,19 +35,13 @@ class HueEvent(GenericHueDevice): self._last_updated = self.sensor.lastupdated # Register callback in coordinator and add job to remove it on bridge reset. - self.bridge.sensor_manager.coordinator.async_add_listener( - self.async_update_callback + self.bridge.reset_jobs.append( + self.bridge.sensor_manager.coordinator.async_add_listener( + self.async_update_callback + ) ) - self.bridge.reset_jobs.append(self.async_will_remove_from_hass) _LOGGER.debug("Hue event created: %s", self.event_id) - @callback - def async_will_remove_from_hass(self): - """Remove listener on bridge reset.""" - self.bridge.sensor_manager.coordinator.async_remove_listener( - self.async_update_callback - ) - @callback def async_update_callback(self): """Fire the event if reason is that state is updated.""" diff --git a/homeassistant/components/hue/light.py b/homeassistant/components/hue/light.py index e468c516676..649be55e94e 100644 --- a/homeassistant/components/hue/light.py +++ b/homeassistant/components/hue/light.py @@ -118,13 +118,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ) # We add a listener after fetching the data, so manually trigger listener - light_coordinator.async_add_listener(update_lights) + bridge.reset_jobs.append(light_coordinator.async_add_listener(update_lights)) update_lights() - bridge.reset_jobs.append( - lambda: light_coordinator.async_remove_listener(update_lights) - ) - api_version = tuple(int(v) for v in bridge.api.config.apiversion.split(".")) allow_groups = bridge.allow_groups @@ -155,13 +151,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities): partial(create_light, HueLight, group_coordinator, bridge, True), ) - group_coordinator.async_add_listener(update_groups) + bridge.reset_jobs.append(group_coordinator.async_add_listener(update_groups)) await group_coordinator.async_refresh() - bridge.reset_jobs.append( - lambda: group_coordinator.async_remove_listener(update_groups) - ) - async def async_safe_fetch(bridge, fetch_method): """Safely fetch data.""" @@ -339,11 +331,9 @@ class HueLight(Light): async def async_added_to_hass(self): """When entity is added to hass.""" - self.coordinator.async_add_listener(self.async_write_ha_state) - - async def async_will_remove_from_hass(self): - """When entity will be removed from hass.""" - self.coordinator.async_remove_listener(self.async_write_ha_state) + self.async_on_remove( + self.coordinator.async_add_listener(self.async_write_ha_state) + ) async def async_turn_on(self, **kwargs): """Turn the specified or all lights on.""" diff --git a/homeassistant/components/hue/sensor_base.py b/homeassistant/components/hue/sensor_base.py index 113957d140e..93b98a7c9ce 100644 --- a/homeassistant/components/hue/sensor_base.py +++ b/homeassistant/components/hue/sensor_base.py @@ -76,9 +76,8 @@ class SensorManager: return # We have all components available, start the updating. - self.coordinator.async_add_listener(self.async_update_items) self.bridge.reset_jobs.append( - lambda: self.coordinator.async_remove_listener(self.async_update_items) + self.coordinator.async_add_listener(self.async_update_items) ) await self.coordinator.async_refresh() @@ -178,14 +177,10 @@ class GenericHueSensor(GenericHueDevice, entity.Entity): async def async_added_to_hass(self): """When entity is added to hass.""" - self.bridge.sensor_manager.coordinator.async_add_listener( - self.async_write_ha_state - ) - - async def async_will_remove_from_hass(self): - """When entity will be removed from hass.""" - self.bridge.sensor_manager.coordinator.async_remove_listener( - self.async_write_ha_state + self.async_on_remove( + self.bridge.sensor_manager.coordinator.async_add_listener( + self.async_write_ha_state + ) ) async def async_update(self): diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py index b2b04816616..d9a79d6555c 100644 --- a/homeassistant/helpers/update_coordinator.py +++ b/homeassistant/helpers/update_coordinator.py @@ -62,7 +62,7 @@ class DataUpdateCoordinator: self._debounced_refresh = request_refresh_debouncer @callback - def async_add_listener(self, update_callback: CALLBACK_TYPE) -> None: + def async_add_listener(self, update_callback: CALLBACK_TYPE) -> Callable[[], None]: """Listen for data updates.""" schedule_refresh = not self._listeners @@ -72,6 +72,13 @@ class DataUpdateCoordinator: if schedule_refresh: self._schedule_refresh() + @callback + def remove_listener() -> None: + """Remove update listener.""" + self.async_remove_listener(update_callback) + + return remove_listener + @callback def async_remove_listener(self, update_callback: CALLBACK_TYPE) -> None: """Remove data update.""" diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py index c17c79ccbc8..17caa1f7478 100644 --- a/tests/helpers/test_update_coordinator.py +++ b/tests/helpers/test_update_coordinator.py @@ -18,11 +18,12 @@ LOGGER = logging.getLogger(__name__) @pytest.fixture def crd(hass): """Coordinator mock.""" - calls = [] + calls = 0 async def refresh(): - calls.append(None) - return len(calls) + nonlocal calls + calls += 1 + return calls crd = update_coordinator.DataUpdateCoordinator( hass, @@ -46,16 +47,19 @@ async def test_async_refresh(crd): def update_callback(): updates.append(crd.data) - crd.async_add_listener(update_callback) - + unsub = crd.async_add_listener(update_callback) await crd.async_refresh() - assert updates == [2] - crd.async_remove_listener(update_callback) - + # Test unsubscribing through function + unsub() await crd.async_refresh() + assert updates == [2] + # Test unsubscribing through method + crd.async_add_listener(update_callback) + crd.async_remove_listener(update_callback) + await crd.async_refresh() assert updates == [2]