Clean up custom polling in ZHA device and light (#32653)

* cleanup timer handle when device is removed

* separate unavailable times for mains vs battery

* better name

* remove light refresh handle when removing light

* remove unused parallel updates

* don't steal HA const for different purpose

* don't flood network every hour for lights

* update test to test both intervals

* add test for light refresh interval

* fix tests

* update test

* put parallel updates back for now...

* fix async_schedule_update_ha_state usage

* review comment

* review comment

* update test - review conversation

* review comments

* await count not call count

* flip some state

homeassistant/components/zha/binary_sensor.py

@@ -106,7 +106,7 @@ class BinarySensor(ZhaEntity, BinarySensorDevice):
     def async_set_state(self, attr_id, attr_name, value):
         """Set the state."""
         self._state = bool(value)
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_update(self):
         """Attempt to retrieve on off state from the binary sensor."""

homeassistant/components/zha/core/device.py

@@ -61,7 +61,8 @@ from .const import (
 from .helpers import LogMixin
 
 _LOGGER = logging.getLogger(__name__)
-_KEEP_ALIVE_INTERVAL = 7200
+_CONSIDER_UNAVAILABLE_MAINS = 60 * 60 * 2  # 2 hours
+_CONSIDER_UNAVAILABLE_BATTERY = 60 * 60 * 6  # 6 hours
 _UPDATE_ALIVE_INTERVAL = (60, 90)
 _CHECKIN_GRACE_PERIODS = 2
 
@@ -99,8 +100,12 @@ class ZHADevice(LogMixin):
             self._zigpy_device.__class__.__module__,
             self._zigpy_device.__class__.__name__,
         )
+        if self.is_mains_powered:
+            self._consider_unavailable_time = _CONSIDER_UNAVAILABLE_MAINS
+        else:
+            self._consider_unavailable_time = _CONSIDER_UNAVAILABLE_BATTERY
         keep_alive_interval = random.randint(*_UPDATE_ALIVE_INTERVAL)
-        self._available_check = async_track_time_interval(
+        self._cancel_available_check = async_track_time_interval(
             self.hass, self._check_available, timedelta(seconds=keep_alive_interval)
         )
         self._ha_device_id = None
@@ -279,7 +284,7 @@ class ZHADevice(LogMixin):
             return
 
         difference = time.time() - self.last_seen
-        if difference < _KEEP_ALIVE_INTERVAL:
+        if difference < self._consider_unavailable_time:
             self.update_available(True)
             self._checkins_missed_count = 0
             return
@@ -363,9 +368,10 @@ class ZHADevice(LogMixin):
         self.debug("completed initialization")
 
     @callback
-    def async_unsub_dispatcher(self):
-        """Unsubscribe the dispatcher."""
+    def async_cleanup_handles(self) -> None:
+        """Unsubscribe the dispatchers and timers."""
         self._unsub()
+        self._cancel_available_check()
 
     @callback
     def async_update_last_seen(self, last_seen):

homeassistant/components/zha/core/gateway.py

@@ -262,7 +262,7 @@ class ZHAGateway:
         entity_refs = self._device_registry.pop(device.ieee, None)
         if zha_device is not None:
             device_info = zha_device.async_get_info()
-            zha_device.async_unsub_dispatcher()
+            zha_device.async_cleanup_handles()
             async_dispatcher_send(
                 self._hass, "{}_{}".format(SIGNAL_REMOVE, str(zha_device.ieee))
             )

homeassistant/components/zha/cover.py

@@ -99,14 +99,14 @@ class ZhaCover(ZhaEntity, CoverDevice):
             self._state = STATE_CLOSED
         elif self._current_position == 100:
             self._state = STATE_OPEN
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @callback
     def async_update_state(self, state):
         """Handle state update from channel."""
         _LOGGER.debug("state=%s", state)
         self._state = state
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_open_cover(self, **kwargs):
         """Open the window cover."""
@@ -134,7 +134,7 @@ class ZhaCover(ZhaEntity, CoverDevice):
         res = await self._cover_channel.stop()
         if isinstance(res, list) and res[1] is Status.SUCCESS:
             self._state = STATE_OPEN if self._current_position > 0 else STATE_CLOSED
-            self.async_schedule_update_ha_state()
+            self.async_write_ha_state()
 
     async def async_update(self):
         """Attempt to retrieve the open/close state of the cover."""

homeassistant/components/zha/device_tracker.py

@@ -90,7 +90,7 @@ class ZHADeviceScannerEntity(ScannerEntity, ZhaEntity):
         self.debug("battery_percentage_remaining updated: %s", value)
         self._connected = True
         self._battery_level = Battery.formatter(value)
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @property
     def battery_level(self):

homeassistant/components/zha/entity.py

@@ -102,13 +102,13 @@ class ZhaEntity(RestoreEntity, LogMixin, entity.Entity):
     def async_set_available(self, available):
         """Set entity availability."""
         self._available = available
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @callback
     def async_update_state_attribute(self, key, value):
         """Update a single device state attribute."""
         self._device_state_attributes.update({key: value})
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @callback
     def async_set_state(self, attr_id, attr_name, value):

homeassistant/components/zha/fan.py

@@ -117,7 +117,7 @@ class ZhaFan(ZhaEntity, FanEntity):
     def async_set_state(self, attr_id, attr_name, value):
         """Handle state update from channel."""
         self._state = VALUE_TO_SPEED.get(value, self._state)
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_turn_on(self, speed: str = None, **kwargs) -> None:
         """Turn the entity on."""

homeassistant/components/zha/light.py

@@ -2,6 +2,7 @@
 from datetime import timedelta
 import functools
 import logging
+import random
 
 from zigpy.zcl.foundation import Status
 
@@ -44,9 +45,9 @@ UPDATE_COLORLOOP_HUE = 0x8
 FLASH_EFFECTS = {light.FLASH_SHORT: EFFECT_BLINK, light.FLASH_LONG: EFFECT_BREATHE}
 
 UNSUPPORTED_ATTRIBUTE = 0x86
-SCAN_INTERVAL = timedelta(minutes=60)
 STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, light.DOMAIN)
-PARALLEL_UPDATES = 5
+PARALLEL_UPDATES = 0
+_REFRESH_INTERVAL = (45, 75)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
@@ -81,6 +82,7 @@ class Light(ZhaEntity, light.Light):
         self._level_channel = self.cluster_channels.get(CHANNEL_LEVEL)
         self._color_channel = self.cluster_channels.get(CHANNEL_COLOR)
         self._identify_channel = self.zha_device.channels.identify_ch
+        self._cancel_refresh_handle = None
 
         if self._level_channel:
             self._supported_features |= light.SUPPORT_BRIGHTNESS
@@ -130,7 +132,7 @@ class Light(ZhaEntity, light.Light):
         """
         value = max(0, min(254, value))
         self._brightness = value
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @property
     def hs_color(self):
@@ -163,7 +165,7 @@ class Light(ZhaEntity, light.Light):
         self._state = bool(value)
         if value:
             self._off_brightness = None
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_added_to_hass(self):
         """Run when about to be added to hass."""
@@ -175,7 +177,15 @@ class Light(ZhaEntity, light.Light):
             await self.async_accept_signal(
                 self._level_channel, SIGNAL_SET_LEVEL, self.set_level
             )
-        async_track_time_interval(self.hass, self.refresh, SCAN_INTERVAL)
+        refresh_interval = random.randint(*_REFRESH_INTERVAL)
+        self._cancel_refresh_handle = async_track_time_interval(
+            self.hass, self._refresh, timedelta(minutes=refresh_interval)
+        )
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Disconnect entity object when removed."""
+        self._cancel_refresh_handle()
+        await super().async_will_remove_from_hass()
 
     @callback
     def async_restore_last_state(self, last_state):
@@ -296,7 +306,7 @@ class Light(ZhaEntity, light.Light):
 
         self._off_brightness = None
         self.debug("turned on: %s", t_log)
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_turn_off(self, **kwargs):
         """Turn the entity off."""
@@ -318,7 +328,7 @@ class Light(ZhaEntity, light.Light):
             # store current brightness so that the next turn_on uses it.
             self._off_brightness = self._brightness
 
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_update(self):
         """Attempt to retrieve on off state from the light."""
@@ -384,6 +394,7 @@ class Light(ZhaEntity, light.Light):
                 if color_loop_active == 1:
                     self._effect = light.EFFECT_COLORLOOP
 
-    async def refresh(self, time):
+    async def _refresh(self, time):
         """Call async_get_state at an interval."""
         await self.async_get_state(from_cache=False)
+        self.async_write_ha_state()

homeassistant/components/zha/lock.py

@@ -87,7 +87,7 @@ class ZhaDoorLock(ZhaEntity, LockDevice):
         if not isinstance(result, list) or result[0] is not Status.SUCCESS:
             self.error("Error with lock_door: %s", result)
             return
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_unlock(self, **kwargs):
         """Unlock the lock."""
@@ -95,7 +95,7 @@ class ZhaDoorLock(ZhaEntity, LockDevice):
         if not isinstance(result, list) or result[0] is not Status.SUCCESS:
             self.error("Error with unlock_door: %s", result)
             return
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_update(self):
         """Attempt to retrieve state from the lock."""
@@ -106,7 +106,7 @@ class ZhaDoorLock(ZhaEntity, LockDevice):
     def async_set_state(self, attr_id, attr_name, value):
         """Handle state update from channel."""
         self._state = VALUE_TO_STATE.get(value, self._state)
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_get_state(self, from_cache=True):
         """Attempt to retrieve state from the lock."""

homeassistant/components/zha/sensor.py

@@ -129,7 +129,7 @@ class Sensor(ZhaEntity):
         if value is not None:
             value = self.formatter(value)
         self._state = value
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @callback
     def async_restore_last_state(self, last_state):
@@ -191,7 +191,7 @@ class Battery(Sensor):
         """Update a single device state attribute."""
         if key == "battery_voltage":
             self._device_state_attributes[key] = round(value / 10, 1)
-            self.async_schedule_update_ha_state()
+            self.async_write_ha_state()
 
 
 @STRICT_MATCH(channel_names=CHANNEL_ELECTRICAL_MEASUREMENT)

homeassistant/components/zha/switch.py

@@ -60,7 +60,7 @@ class Switch(ZhaEntity, SwitchDevice):
         if not isinstance(result, list) or result[1] is not Status.SUCCESS:
             return
         self._state = True
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     async def async_turn_off(self, **kwargs):
         """Turn the entity off."""
@@ -68,13 +68,13 @@ class Switch(ZhaEntity, SwitchDevice):
         if not isinstance(result, list) or result[1] is not Status.SUCCESS:
             return
         self._state = False
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @callback
     def async_set_state(self, attr_id, attr_name, value):
         """Handle state update from channel."""
         self._state = bool(value)
-        self.async_schedule_update_ha_state()
+        self.async_write_ha_state()
 
     @property
     def device_state_attributes(self):

tests/components/zha/test_device.py

@@ -8,11 +8,12 @@ import pytest
 import zigpy.zcl.clusters.general as general
 
 import homeassistant.components.zha.core.device as zha_core_device
-import homeassistant.core as ha
 import homeassistant.util.dt as dt_util
 
 from .common import async_enable_traffic
 
+from tests.common import async_fire_time_changed
+
 
 @pytest.fixture
 def zigpy_device(zigpy_device_mock):
@@ -32,9 +33,28 @@ def zigpy_device(zigpy_device_mock):
 
 
 @pytest.fixture
-def device_with_basic_channel(zigpy_device):
+def zigpy_device_mains(zigpy_device_mock):
+    """Device tracker zigpy device."""
+
+    def _dev(with_basic_channel: bool = True):
+        in_clusters = [general.OnOff.cluster_id]
+        if with_basic_channel:
+            in_clusters.append(general.Basic.cluster_id)
+
+        endpoints = {
+            3: {"in_clusters": in_clusters, "out_clusters": [], "device_type": 0}
+        }
+        return zigpy_device_mock(
+            endpoints, node_descriptor=b"\x02@\x84_\x11\x7fd\x00\x00,d\x00\x00"
+        )
+
+    return _dev
+
+
+@pytest.fixture
+def device_with_basic_channel(zigpy_device_mains):
     """Return a zha device with a basic channel present."""
-    return zigpy_device(with_basic_channel=True)
+    return zigpy_device_mains(with_basic_channel=True)
 
 
 @pytest.fixture
@@ -45,8 +65,8 @@ def device_without_basic_channel(zigpy_device):
 
 def _send_time_changed(hass, seconds):
     """Send a time changed event."""
-    now = dt_util.utcnow() + timedelta(seconds)
-    hass.bus.async_fire(ha.EVENT_TIME_CHANGED, {ha.ATTR_NOW: now})
+    now = dt_util.utcnow() + timedelta(seconds=seconds)
+    async_fire_time_changed(hass, now)
 
 
 @asynctest.patch(
@@ -66,13 +86,13 @@ async def test_check_available_success(
     basic_ch.read_attributes.reset_mock()
     device_with_basic_channel.last_seen = None
     assert zha_device.available is True
-    _send_time_changed(hass, 61)
+    _send_time_changed(hass, zha_core_device._CONSIDER_UNAVAILABLE_MAINS + 2)
     await hass.async_block_till_done()
     assert zha_device.available is False
     assert basic_ch.read_attributes.await_count == 0
 
     device_with_basic_channel.last_seen = (
-        time.time() - zha_core_device._KEEP_ALIVE_INTERVAL - 2
+        time.time() - zha_core_device._CONSIDER_UNAVAILABLE_MAINS - 2
     )
     _seens = [time.time(), device_with_basic_channel.last_seen]
 
@@ -121,7 +141,7 @@ async def test_check_available_unsuccessful(
     assert basic_ch.read_attributes.await_count == 0
 
     device_with_basic_channel.last_seen = (
-        time.time() - zha_core_device._KEEP_ALIVE_INTERVAL - 2
+        time.time() - zha_core_device._CONSIDER_UNAVAILABLE_MAINS - 2
     )
 
     # unsuccessfuly ping zigpy device, but zha_device is still available
@@ -162,7 +182,7 @@ async def test_check_available_no_basic_channel(
     assert zha_device.available is True
 
     device_without_basic_channel.last_seen = (
-        time.time() - zha_core_device._KEEP_ALIVE_INTERVAL - 2
+        time.time() - zha_core_device._CONSIDER_UNAVAILABLE_BATTERY - 2
     )
 
     assert "does not have a mandatory basic cluster" not in caplog.text

tests/components/zha/test_light.py

@@ -1,5 +1,6 @@
 """Test zha light."""
-from unittest.mock import call, sentinel
+from datetime import timedelta
+from unittest.mock import MagicMock, call, sentinel
 
 import asynctest
 import pytest
@@ -12,6 +13,7 @@ import zigpy.zcl.foundation as zcl_f
 from homeassistant.components.light import DOMAIN, FLASH_LONG, FLASH_SHORT
 from homeassistant.components.zha.light import FLASH_EFFECTS
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
+import homeassistant.util.dt as dt_util
 
 from .common import (
     async_enable_traffic,
@@ -21,6 +23,8 @@ from .common import (
     make_zcl_header,
 )
 
+from tests.common import async_fire_time_changed
+
 ON = 1
 OFF = 0
 
@@ -63,6 +67,46 @@ LIGHT_COLOR = {
 }
 
 
+@asynctest.mock.patch(
+    "zigpy.zcl.clusters.general.OnOff.read_attributes", new=MagicMock()
+)
+async def test_light_refresh(hass, zigpy_device_mock, zha_device_joined_restored):
+    """Test zha light platform refresh."""
+
+    # create zigpy devices
+    zigpy_device = zigpy_device_mock(LIGHT_ON_OFF)
+    zha_device = await zha_device_joined_restored(zigpy_device)
+    on_off_cluster = zigpy_device.endpoints[1].on_off
+    entity_id = await find_entity_id(DOMAIN, zha_device, hass)
+
+    # allow traffic to flow through the gateway and device
+    await async_enable_traffic(hass, [zha_device])
+    on_off_cluster.read_attributes.reset_mock()
+
+    # not enough time passed
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=20))
+    await hass.async_block_till_done()
+    assert on_off_cluster.read_attributes.call_count == 0
+    assert on_off_cluster.read_attributes.await_count == 0
+    assert hass.states.get(entity_id).state == STATE_OFF
+
+    # 1 interval - 1 call
+    on_off_cluster.read_attributes.return_value = [{"on_off": 1}, {}]
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=80))
+    await hass.async_block_till_done()
+    assert on_off_cluster.read_attributes.call_count == 1
+    assert on_off_cluster.read_attributes.await_count == 1
+    assert hass.states.get(entity_id).state == STATE_ON
+
+    # 2 intervals - 2 calls
+    on_off_cluster.read_attributes.return_value = [{"on_off": 0}, {}]
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=80))
+    await hass.async_block_till_done()
+    assert on_off_cluster.read_attributes.call_count == 2
+    assert on_off_cluster.read_attributes.await_count == 2
+    assert hass.states.get(entity_id).state == STATE_OFF
+
+
 @asynctest.patch(
     "zigpy.zcl.clusters.lighting.Color.request",
     new=asynctest.CoroutineMock(return_value=[sentinel.data, zcl_f.Status.SUCCESS]),
@@ -84,7 +128,7 @@ LIGHT_COLOR = {
     [(LIGHT_ON_OFF, (1, 0, 0)), (LIGHT_LEVEL, (1, 1, 0)), (LIGHT_COLOR, (1, 1, 3))],
 )
 async def test_light(
-    hass, zigpy_device_mock, zha_device_joined_restored, device, reporting,
+    hass, zigpy_device_mock, zha_device_joined_restored, device, reporting
 ):
     """Test zha light platform."""