Cleanup inconsistencies in zha fan and make it a bit more dry (#46714)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/zha/fan.py

@@ -13,11 +13,13 @@ from homeassistant.components.fan import (
     DOMAIN,
     SUPPORT_SET_SPEED,
     FanEntity,
+    NotValidPresetModeError,
 )
 from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import State, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.util.percentage import (
+    int_states_in_range,
     percentage_to_ranged_value,
     ranged_value_to_percentage,
 )
@@ -75,7 +77,7 @@ class BaseFan(FanEntity):
     """Base representation of a ZHA fan."""
 
     @property
-    def preset_modes(self) -> str:
+    def preset_modes(self) -> List[str]:
         """Return the available preset modes."""
         return PRESET_MODES
 
@@ -84,10 +86,17 @@ class BaseFan(FanEntity):
         """Flag supported features."""
         return SUPPORT_SET_SPEED
 
+    @property
+    def speed_count(self) -> int:
+        """Return the number of speeds the fan supports."""
+        return int_states_in_range(SPEED_RANGE)
+
     async def async_turn_on(
         self, speed=None, percentage=None, preset_mode=None, **kwargs
     ) -> None:
         """Turn the entity on."""
+        if percentage is None:
+            percentage = DEFAULT_ON_PERCENTAGE
         await self.async_set_percentage(percentage)
 
     async def async_turn_off(self, **kwargs) -> None:
@@ -96,15 +105,16 @@ class BaseFan(FanEntity):
 
     async def async_set_percentage(self, percentage: Optional[int]) -> None:
         """Set the speed percenage of the fan."""
-        if percentage is None:
-            percentage = DEFAULT_ON_PERCENTAGE
         fan_mode = math.ceil(percentage_to_ranged_value(SPEED_RANGE, percentage))
         await self._async_set_fan_mode(fan_mode)
 
     async def async_set_preset_mode(self, preset_mode: str) -> None:
-        """Set the speed percenage of the fan."""
-        fan_mode = NAME_TO_PRESET_MODE.get(preset_mode)
-        await self._async_set_fan_mode(fan_mode)
+        """Set the preset mode for the fan."""
+        if preset_mode not in self.preset_modes:
+            raise NotValidPresetModeError(
+                f"The preset_mode {preset_mode} is not a valid preset_mode: {self.preset_modes}"
+            )
+        await self._async_set_fan_mode(NAME_TO_PRESET_MODE[preset_mode])
 
     @abstractmethod
     async def _async_set_fan_mode(self, fan_mode: int) -> None:
@@ -132,7 +142,7 @@ class ZhaFan(BaseFan, ZhaEntity):
         )
 
     @property
-    def percentage(self) -> str:
+    def percentage(self) -> Optional[int]:
         """Return the current speed percentage."""
         if (
             self._fan_channel.fan_mode is None
@@ -144,7 +154,7 @@ class ZhaFan(BaseFan, ZhaEntity):
         return ranged_value_to_percentage(SPEED_RANGE, self._fan_channel.fan_mode)
 
     @property
-    def preset_mode(self) -> str:
+    def preset_mode(self) -> Optional[str]:
         """Return the current preset mode."""
         return PRESET_MODES_TO_NAME.get(self._fan_channel.fan_mode)
 
@@ -175,27 +185,15 @@ class FanGroup(BaseFan, ZhaGroupEntity):
         self._preset_mode = None
 
     @property
-    def percentage(self) -> str:
+    def percentage(self) -> Optional[int]:
         """Return the current speed percentage."""
         return self._percentage
 
     @property
-    def preset_mode(self) -> str:
+    def preset_mode(self) -> Optional[str]:
         """Return the current preset mode."""
         return self._preset_mode
 
-    async def async_set_percentage(self, percentage: Optional[int]) -> None:
-        """Set the speed percenage of the fan."""
-        if percentage is None:
-            percentage = DEFAULT_ON_PERCENTAGE
-        fan_mode = math.ceil(percentage_to_ranged_value(SPEED_RANGE, percentage))
-        await self._async_set_fan_mode(fan_mode)
-
-    async def async_set_preset_mode(self, preset_mode: str) -> None:
-        """Set the speed percenage of the fan."""
-        fan_mode = NAME_TO_PRESET_MODE.get(preset_mode)
-        await self._async_set_fan_mode(fan_mode)
-
     async def _async_set_fan_mode(self, fan_mode: int) -> None:
         """Set the fan mode for the group."""
         try:

tests/components/zha/test_fan.py

@@ -11,6 +11,7 @@ import zigpy.zcl.foundation as zcl_f
 from homeassistant.components import fan
 from homeassistant.components.fan import (
     ATTR_PERCENTAGE,
+    ATTR_PERCENTAGE_STEP,
     ATTR_PRESET_MODE,
     ATTR_SPEED,
     DOMAIN,
@@ -20,6 +21,7 @@ from homeassistant.components.fan import (
     SPEED_LOW,
     SPEED_MEDIUM,
     SPEED_OFF,
+    NotValidPresetModeError,
 )
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.components.zha.core.discovery import GROUP_PROBE
@@ -188,6 +190,14 @@ async def test_fan(hass, zha_device_joined_restored, zigpy_device):
     assert len(cluster.write_attributes.mock_calls) == 1
     assert cluster.write_attributes.call_args == call({"fan_mode": 4})
 
+    # set invalid preset_mode from HA
+    cluster.write_attributes.reset_mock()
+    with pytest.raises(NotValidPresetModeError):
+        await async_set_preset_mode(
+            hass, entity_id, preset_mode="invalid does not exist"
+        )
+    assert len(cluster.write_attributes.mock_calls) == 0
+
     # test adding new fan to the network and HA
     await async_test_rejoin(hass, zigpy_device, [cluster], (1,))
 
@@ -450,6 +460,7 @@ async def test_fan_update_entity(
     assert hass.states.get(entity_id).attributes[ATTR_SPEED] == SPEED_OFF
     assert hass.states.get(entity_id).attributes[ATTR_PERCENTAGE] == 0
     assert hass.states.get(entity_id).attributes[ATTR_PRESET_MODE] is None
+    assert hass.states.get(entity_id).attributes[ATTR_PERCENTAGE_STEP] == 100 / 3
     assert cluster.read_attributes.await_count == 1
 
     await async_setup_component(hass, "homeassistant", {})
@@ -470,4 +481,5 @@ async def test_fan_update_entity(
     assert hass.states.get(entity_id).attributes[ATTR_PERCENTAGE] == 33
     assert hass.states.get(entity_id).attributes[ATTR_SPEED] == SPEED_LOW
     assert hass.states.get(entity_id).attributes[ATTR_PRESET_MODE] is None
+    assert hass.states.get(entity_id).attributes[ATTR_PERCENTAGE_STEP] == 100 / 3
     assert cluster.read_attributes.await_count == 3