Data Coordinator to return unsub func

homeassistant/components/hue/hue_event.py

@@ -35,19 +35,13 @@ def __init__(self, sensor, name, bridge, primary_sensor=None):
         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."""

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 @@ def device_info(self):
 
     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."""

homeassistant/components/hue/sensor_base.py

@@ -76,9 +76,8 @@ async def async_register_component(self, platform, async_add_entities):
             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 @@ def available(self):
 
     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):

homeassistant/helpers/update_coordinator.py

@@ -62,7 +62,7 @@ def __init__(
         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 @@ def async_add_listener(self, update_callback: CALLBACK_TYPE) -> None:
         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."""

tests/helpers/test_update_coordinator.py

@@ -18,11 +18,12 @@
 @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]
 
+    # 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]