Cleanup ZHA fan channel (#43973)

* Use zigpy cached values for ZHA Fan speed

* Disable update_before_add for ZHA fans

* Refresh state of the group

* Fix group tests

homeassistant/components/zha/core/channels/hvac.py

@@ -43,17 +43,10 @@ class FanChannel(ZigbeeChannel):
 
     REPORT_CONFIG = ({"attr": "fan_mode", "config": REPORT_CONFIG_OP},)
 
-    def __init__(
-        self, cluster: zha_typing.ZigpyClusterType, ch_pool: zha_typing.ChannelPoolType
-    ):
-        """Init Thermostat channel instance."""
-        super().__init__(cluster, ch_pool)
-        self._fan_mode = None
-
     @property
     def fan_mode(self) -> Optional[int]:
         """Return current fan mode."""
-        return self._fan_mode
+        return self.cluster.get("fan_mode")
 
     async def async_set_speed(self, value) -> None:
         """Set the speed of the fan."""
@@ -66,12 +59,7 @@ class FanChannel(ZigbeeChannel):
 
     async def async_update(self) -> None:
         """Retrieve latest state."""
-        result = await self.get_attribute_value("fan_mode", from_cache=True)
+        await self.get_attribute_value("fan_mode", from_cache=False)
-        if result is not None:
-            self._fan_mode = result
-            self.async_send_signal(
-                f"{self.unique_id}_{SIGNAL_ATTR_UPDATED}", 0, "fan_mode", result
-            )
 
     @callback
     def attribute_updated(self, attrid: int, value: Any) -> None:
@@ -80,8 +68,7 @@ class FanChannel(ZigbeeChannel):
         self.debug(
             "Attribute report '%s'[%s] = %s", self.cluster.name, attr_name, value
         )
-        if attrid == self._value_attribute:
+        if attr_name == "fan_mode":
-            self._fan_mode = value
             self.async_send_signal(
                 f"{self.unique_id}_{SIGNAL_ATTR_UPDATED}", attrid, attr_name, value
             )

homeassistant/components/zha/fan.py

@@ -1,6 +1,6 @@
 """Fans on Zigbee Home Automation networks."""
 import functools
-from typing import List
+from typing import List, Optional
 
 from zigpy.exceptions import ZigbeeException
 import zigpy.zcl.clusters.hvac as hvac
@@ -62,7 +62,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         hass,
         SIGNAL_ADD_ENTITIES,
         functools.partial(
-            discovery.async_add_entities, async_add_entities, entities_to_create
+            discovery.async_add_entities,
+            async_add_entities,
+            entities_to_create,
+            update_before_add=False,
         ),
     )
     hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS].append(unsub)
@@ -87,13 +90,6 @@ class BaseFan(FanEntity):
         """Return the current speed."""
         return self._state
 
-    @property
-    def is_on(self) -> bool:
-        """Return true if entity is on."""
-        if self._state is None:
-            return False
-        return self._state != SPEED_OFF
-
     @property
     def supported_features(self) -> int:
         """Flag supported features."""
@@ -136,25 +132,16 @@ class ZhaFan(BaseFan, ZhaEntity):
             self._fan_channel, SIGNAL_ATTR_UPDATED, self.async_set_state
         )
 
-    @callback
+    @property
-    def async_restore_last_state(self, last_state):
+    def speed(self) -> Optional[str]:
-        """Restore previous state."""
+        """Return the current speed."""
-        self._state = VALUE_TO_SPEED.get(last_state.state, last_state.state)
+        return VALUE_TO_SPEED.get(self._fan_channel.fan_mode)
 
     @callback
     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_write_ha_state()
 
-    async def async_update(self):
-        """Attempt to retrieve on off state from the fan."""
-        await super().async_update()
-        if self._fan_channel:
-            state = await self._fan_channel.get_attribute_value("fan_mode")
-            if state is not None:
-                self._state = VALUE_TO_SPEED.get(state, self._state)
-
 
 @GROUP_MATCH()
 class FanGroup(BaseFan, ZhaGroupEntity):
@@ -185,9 +172,15 @@ class FanGroup(BaseFan, ZhaGroupEntity):
         all_states = [self.hass.states.get(x) for x in self._entity_ids]
         states: List[State] = list(filter(None, all_states))
         on_states: List[State] = [state for state in states if state.state != SPEED_OFF]
+
         self._available = any(state.state != STATE_UNAVAILABLE for state in states)
         # for now just use first non off state since its kind of arbitrary
         if not on_states:
             self._state = SPEED_OFF
         else:
-            self._state = states[0].state
+            self._state = on_states[0].state
+
+    async def async_added_to_hass(self) -> None:
+        """Run when about to be added to hass."""
+        await self.async_update()
+        await super().async_added_to_hass()

tests/components/zha/test_fan.py

@@ -3,6 +3,7 @@ import pytest
 import zigpy.profiles.zha as zha
 import zigpy.zcl.clusters.general as general
 import zigpy.zcl.clusters.hvac as hvac
+import zigpy.zcl.foundation as zcl_f
 
 from homeassistant.components import fan
 from homeassistant.components.fan import (
@@ -10,11 +11,13 @@ from homeassistant.components.fan import (
     DOMAIN,
     SERVICE_SET_SPEED,
     SPEED_HIGH,
+    SPEED_LOW,
     SPEED_MEDIUM,
     SPEED_OFF,
 )
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.components.zha.core.discovery import GROUP_PROBE
+from homeassistant.components.zha.core.group import GroupMember
 from homeassistant.const import (
     ATTR_ENTITY_ID,
     SERVICE_TURN_OFF,
@@ -23,6 +26,7 @@ from homeassistant.const import (
     STATE_ON,
     STATE_UNAVAILABLE,
 )
+from homeassistant.setup import async_setup_component
 
 from .common import (
     async_enable_traffic,
@@ -33,7 +37,7 @@ from .common import (
     send_attributes_report,
 )
 
-from tests.async_mock import call
+from tests.async_mock import AsyncMock, call, patch
 
 IEEE_GROUPABLE_DEVICE = "01:2d:6f:00:0a:90:69:e8"
 IEEE_GROUPABLE_DEVICE2 = "02:2d:6f:00:0a:90:69:e8"
@@ -49,7 +53,9 @@ def zigpy_device(zigpy_device_mock):
             "device_type": zha.DeviceType.ON_OFF_SWITCH,
         }
     }
-    return zigpy_device_mock(endpoints)
+    return zigpy_device_mock(
+        endpoints, node_descriptor=b"\x02@\x8c\x02\x10RR\x00\x00\x00R\x00\x00"
+    )
 
 
 @pytest.fixture
@@ -59,7 +65,7 @@ async def coordinator(hass, zigpy_device_mock, zha_device_joined):
     zigpy_device = zigpy_device_mock(
         {
             1: {
-                "in_clusters": [],
+                "in_clusters": [general.Groups.cluster_id],
                 "out_clusters": [],
                 "device_type": zha.DeviceType.COLOR_DIMMABLE_LIGHT,
             }
@@ -80,14 +86,20 @@ async def device_fan_1(hass, zigpy_device_mock, zha_device_joined):
     zigpy_device = zigpy_device_mock(
         {
             1: {
-                "in_clusters": [general.OnOff.cluster_id, hvac.Fan.cluster_id],
+                "in_clusters": [
+                    general.Groups.cluster_id,
+                    general.OnOff.cluster_id,
+                    hvac.Fan.cluster_id,
+                ],
                 "out_clusters": [],
-            }
+                "device_type": zha.DeviceType.ON_OFF_LIGHT,
+            },
         },
         ieee=IEEE_GROUPABLE_DEVICE,
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.available = True
+    await hass.async_block_till_done()
     return zha_device
 
 
@@ -99,17 +111,20 @@ async def device_fan_2(hass, zigpy_device_mock, zha_device_joined):
         {
             1: {
                 "in_clusters": [
+                    general.Groups.cluster_id,
                     general.OnOff.cluster_id,
                     hvac.Fan.cluster_id,
                     general.LevelControl.cluster_id,
                 ],
                 "out_clusters": [],
-            }
+                "device_type": zha.DeviceType.ON_OFF_LIGHT,
+            },
         },
         ieee=IEEE_GROUPABLE_DEVICE2,
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.available = True
+    await hass.async_block_till_done()
     return zha_device
 
 
@@ -191,9 +206,11 @@ async def async_set_speed(hass, entity_id, speed=None):
     await hass.services.async_call(DOMAIN, SERVICE_SET_SPEED, data, blocking=True)
 
 
-async def async_test_zha_group_fan_entity(
+@patch(
-    hass, device_fan_1, device_fan_2, coordinator
+    "zigpy.zcl.clusters.hvac.Fan.write_attributes",
-):
+    new=AsyncMock(return_value=zcl_f.WriteAttributesResponse.deserialize(b"\x00")[0]),
+)
+async def test_zha_group_fan_entity(hass, device_fan_1, device_fan_2, coordinator):
     """Test the fan entity for a ZHA group."""
     zha_gateway = get_zha_gateway(hass)
     assert zha_gateway is not None
@@ -202,19 +219,20 @@ async def async_test_zha_group_fan_entity(
     device_fan_1._zha_gateway = zha_gateway
     device_fan_2._zha_gateway = zha_gateway
     member_ieee_addresses = [device_fan_1.ieee, device_fan_2.ieee]
+    members = [GroupMember(device_fan_1.ieee, 1), GroupMember(device_fan_2.ieee, 1)]
 
     # test creating a group with 2 members
-    zha_group = await zha_gateway.async_create_zigpy_group(
+    zha_group = await zha_gateway.async_create_zigpy_group("Test Group", members)
-        "Test Group", member_ieee_addresses
-    )
     await hass.async_block_till_done()
 
     assert zha_group is not None
     assert len(zha_group.members) == 2
     for member in zha_group.members:
-        assert member.ieee in member_ieee_addresses
+        assert member.device.ieee in member_ieee_addresses
+        assert member.group == zha_group
+        assert member.endpoint is not None
 
-    entity_domains = GROUP_PROBE.determine_entity_domains(zha_group)
+    entity_domains = GROUP_PROBE.determine_entity_domains(hass, zha_group)
     assert len(entity_domains) == 2
 
     assert LIGHT_DOMAIN in entity_domains
@@ -224,14 +242,17 @@ async def async_test_zha_group_fan_entity(
     assert hass.states.get(entity_id) is not None
 
     group_fan_cluster = zha_group.endpoint[hvac.Fan.cluster_id]
-    dev1_fan_cluster = device_fan_1.endpoints[1].fan
-    dev2_fan_cluster = device_fan_2.endpoints[1].fan
 
-    # test that the lights were created and that they are unavailable
+    dev1_fan_cluster = device_fan_1.device.endpoints[1].fan
+    dev2_fan_cluster = device_fan_2.device.endpoints[1].fan
+
+    await async_enable_traffic(hass, [device_fan_1, device_fan_2], enabled=False)
+    await hass.async_block_till_done()
+    # test that the fans were created and that they are unavailable
     assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
 
     # allow traffic to flow through the gateway and device
-    await async_enable_traffic(hass, zha_group.members)
+    await async_enable_traffic(hass, [device_fan_1, device_fan_2])
 
     # test that the fan group entity was created and is off
     assert hass.states.get(entity_id).state == STATE_OFF
@@ -239,37 +260,103 @@ async def async_test_zha_group_fan_entity(
     # turn on from HA
     group_fan_cluster.write_attributes.reset_mock()
     await async_turn_on(hass, entity_id)
+    await hass.async_block_till_done()
     assert len(group_fan_cluster.write_attributes.mock_calls) == 1
-    assert group_fan_cluster.write_attributes.call_args == call({"fan_mode": 2})
+    assert group_fan_cluster.write_attributes.call_args[0][0] == {"fan_mode": 2}
-    assert hass.states.get(entity_id).state == SPEED_MEDIUM
 
     # turn off from HA
     group_fan_cluster.write_attributes.reset_mock()
     await async_turn_off(hass, entity_id)
     assert len(group_fan_cluster.write_attributes.mock_calls) == 1
-    assert group_fan_cluster.write_attributes.call_args == call({"fan_mode": 0})
+    assert group_fan_cluster.write_attributes.call_args[0][0] == {"fan_mode": 0}
-    assert hass.states.get(entity_id).state == STATE_OFF
 
     # change speed from HA
     group_fan_cluster.write_attributes.reset_mock()
     await async_set_speed(hass, entity_id, speed=fan.SPEED_HIGH)
     assert len(group_fan_cluster.write_attributes.mock_calls) == 1
-    assert group_fan_cluster.write_attributes.call_args == call({"fan_mode": 3})
+    assert group_fan_cluster.write_attributes.call_args[0][0] == {"fan_mode": 3}
-    assert hass.states.get(entity_id).state == SPEED_HIGH
 
     # test some of the group logic to make sure we key off states correctly
-    await dev1_fan_cluster.async_set_speed(SPEED_OFF)
+    await send_attributes_report(hass, dev1_fan_cluster, {0: 0})
-    await dev2_fan_cluster.async_set_speed(SPEED_OFF)
+    await send_attributes_report(hass, dev2_fan_cluster, {0: 0})
 
     # test that group fan is off
     assert hass.states.get(entity_id).state == STATE_OFF
 
-    await dev1_fan_cluster.async_set_speed(SPEED_MEDIUM)
+    await send_attributes_report(hass, dev2_fan_cluster, {0: 2})
+    await hass.async_block_till_done()
 
     # test that group fan is speed medium
-    assert hass.states.get(entity_id).state == SPEED_MEDIUM
+    assert hass.states.get(entity_id).state == STATE_ON
 
-    await dev1_fan_cluster.async_set_speed(SPEED_OFF)
+    await send_attributes_report(hass, dev2_fan_cluster, {0: 0})
+    await hass.async_block_till_done()
 
     # test that group fan is now off
     assert hass.states.get(entity_id).state == STATE_OFF
+
+
+@pytest.mark.parametrize(
+    "plug_read, expected_state, expected_speed",
+    (
+        (None, STATE_OFF, None),
+        ({"fan_mode": 0}, STATE_OFF, SPEED_OFF),
+        ({"fan_mode": 1}, STATE_ON, SPEED_LOW),
+        ({"fan_mode": 2}, STATE_ON, SPEED_MEDIUM),
+        ({"fan_mode": 3}, STATE_ON, SPEED_HIGH),
+    ),
+)
+async def test_fan_init(
+    hass,
+    zha_device_joined_restored,
+    zigpy_device,
+    plug_read,
+    expected_state,
+    expected_speed,
+):
+    """Test zha fan platform."""
+
+    cluster = zigpy_device.endpoints.get(1).fan
+    cluster.PLUGGED_ATTR_READS = plug_read
+
+    zha_device = await zha_device_joined_restored(zigpy_device)
+    entity_id = await find_entity_id(DOMAIN, zha_device, hass)
+    assert entity_id is not None
+    assert hass.states.get(entity_id).state == expected_state
+    assert hass.states.get(entity_id).attributes[ATTR_SPEED] == expected_speed
+
+
+async def test_fan_update_entity(
+    hass,
+    zha_device_joined_restored,
+    zigpy_device,
+):
+    """Test zha fan platform."""
+
+    cluster = zigpy_device.endpoints.get(1).fan
+    cluster.PLUGGED_ATTR_READS = {"fan_mode": 0}
+
+    zha_device = await zha_device_joined_restored(zigpy_device)
+    entity_id = await find_entity_id(DOMAIN, zha_device, hass)
+    assert entity_id is not None
+    assert hass.states.get(entity_id).state == STATE_OFF
+    assert hass.states.get(entity_id).attributes[ATTR_SPEED] == SPEED_OFF
+    assert cluster.read_attributes.await_count == 1
+
+    await async_setup_component(hass, "homeassistant", {})
+    await hass.async_block_till_done()
+
+    await hass.services.async_call(
+        "homeassistant", "update_entity", {"entity_id": entity_id}, blocking=True
+    )
+    assert hass.states.get(entity_id).state == STATE_OFF
+    assert hass.states.get(entity_id).attributes[ATTR_SPEED] == SPEED_OFF
+    assert cluster.read_attributes.await_count == 2
+
+    cluster.PLUGGED_ATTR_READS = {"fan_mode": 1}
+    await hass.services.async_call(
+        "homeassistant", "update_entity", {"entity_id": entity_id}, blocking=True
+    )
+    assert hass.states.get(entity_id).state == STATE_ON
+    assert hass.states.get(entity_id).attributes[ATTR_SPEED] == SPEED_LOW
+    assert cluster.read_attributes.await_count == 3