Adds work modes specific to model (#128)

* Adds aditional work modes for Pure 500

* Simplify workmode

* Set fan speed 0 if powered off

* Fix backward compatibility

custom_components/wellbeing/api.py

@@ -2,6 +2,7 @@
 
 import logging
 from enum import Enum
+from typing import cast, Any, Callable
 
 from homeassistant.components.binary_sensor import BinarySensorDeviceClass
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
@@ -37,11 +38,22 @@
 _LOGGER: logging.Logger = logging.getLogger(__package__)
 
 
-class Mode(str, Enum):
+class Model(str, Enum):
+    Muju = "Muju"
+    WELLA5 = "WELLA5"
+    WELLA7 = "WELLA7"
+    PUREA9 = "PUREA9"
+    AX5 = "AX5"
+    AX7 = "AX7"
+    AX9 = "AX9"
+
+class WorkMode(str, Enum):
     OFF = "PowerOff"
-    AUTO = "Auto"
     MANUAL = "Manual"
     UNDEFINED = "Undefined"
+    SMART = "Smart"
+    QUITE = "Quiet"
+    AUTO = "Auto"
 
 
 class ApplianceEntity:
@@ -113,12 +125,13 @@ class Appliance:
     brand: str
     device: str
     firmware: str
-    mode: Mode
+    mode: WorkMode
     entities: list
     capabilities: dict
+    model: Model
 
     def __init__(self, name, pnc_id, model) -> None:
-        self.model = model
+        self.model = Model(model)
         self.pnc_id = pnc_id
         self.name = name
 
@@ -259,32 +272,51 @@ def has_capability(self, capability) -> bool:
         return capability in self.capabilities and self.capabilities[capability]["access"] == "readwrite"
 
     def clear_mode(self):
-        self.mode = Mode.UNDEFINED
+        self.mode = WorkMode.UNDEFINED
+
+    def set_mode(self, mode: WorkMode):
+        self.mode = mode
 
     def setup(self, data, capabilities):
         self.firmware = data.get("FrmVer_NIU")
-        self.mode = Mode(data.get("Workmode"))
+        self.mode = WorkMode(data.get("Workmode"))
+
         self.capabilities = capabilities
         self.entities = [entity.setup(data) for entity in Appliance._create_entities(data) if entity.attr in data]
 
+    @property
+    def preset_modes(self) -> list[WorkMode]:
+        if self.model == Model.Muju:
+            return [WorkMode.SMART, WorkMode.QUITE, WorkMode.MANUAL, WorkMode.OFF]
+        return [WorkMode.AUTO, WorkMode.MANUAL, WorkMode.OFF]
+
+    def work_mode_from_preset_mode(self, preset_mode: str | None) -> WorkMode:
+        if preset_mode:
+            return WorkMode(preset_mode)
+        if self.model == Model.Muju:
+            return WorkMode.SMART
+        return WorkMode.AUTO
+
     @property
     def speed_range(self) -> tuple[int, int]:
         ## Electrolux Devices:
-        if self.model == "Muju":
-            return 1, 3
-        if self.model == "WELLA5":
+        if self.model == Model.Muju:
+            if self.mode is WorkMode.QUITE:
+                return 1, 2
+            return 1, 5
+        if self.model == Model.WELLA5:
             return 1, 5
-        if self.model == "WELLA7":
+        if self.model == Model.WELLA7:
             return 1, 5
-        if self.model == "PUREA9":
+        if self.model == Model.PUREA9:
             return 1, 9
 
         ## AEG Devices:
-        if self.model == "AX5":
+        if self.model == Model.AX5:
             return 1, 5
-        if self.model == "AX7":
+        if self.model == Model.AX7:
             return 1, 5
-        if self.model == "AX9":
+        if self.model == Model.AX9:
             return 1, 9
 
         return 0, 0
@@ -349,7 +381,7 @@ async def set_fan_speed(self, pnc_id: str, level: int):
         result = await appliance.send_command(data)
         _LOGGER.debug(f"Set Fan Speed: {result}")
 
-    async def set_work_mode(self, pnc_id: str, mode: Mode):
+    async def set_work_mode(self, pnc_id: str, mode: WorkMode):
         data = {"Workmode": mode.value}
         appliance = self._api_appliances.get(pnc_id, None)
         if appliance is None:

custom_components/wellbeing/fan.py

@@ -7,19 +7,18 @@
 from homeassistant.components.fan import FanEntity, FanEntityFeature
 from homeassistant.const import Platform
 from homeassistant.util.percentage import percentage_to_ranged_value, ranged_value_to_percentage
+
 from . import WellbeingDataUpdateCoordinator
-from .api import Mode
+from .api import WorkMode
 from .const import DOMAIN
 from .entity import WellbeingEntity
 
 _LOGGER: logging.Logger = logging.getLogger(__package__)
 
 SUPPORTED_FEATURES = (
-    FanEntityFeature.SET_SPEED | FanEntityFeature.PRESET_MODE | FanEntityFeature.TURN_OFF | FanEntityFeature.TURN_ON
+    FanEntityFeature.SET_SPEED | FanEntityFeature.PRESET_MODE
 )
 
-PRESET_MODES = [Mode.OFF, Mode.AUTO, Mode.MANUAL]
-
 
 async def async_setup_entry(hass, entry, async_add_devices):
     """Setup sensor platform."""
@@ -40,9 +39,12 @@ async def async_setup_entry(hass, entry, async_add_devices):
 class WellbeingFan(WellbeingEntity, FanEntity):
     """wellbeing Sensor class."""
 
+    # Add FanEntityFeature.TURN_OFF | FanEntityFeature.TURN_ON and set to True before 2025.2
+    _enable_turn_on_off_backwards_compatibility = False
+
     def __init__(self, coordinator: WellbeingDataUpdateCoordinator, config_entry, pnc_id, entity_type, entity_attr):
         super().__init__(coordinator, config_entry, pnc_id, entity_type, entity_attr)
-        self._preset_mode = str(self.get_appliance.mode)
+        self._preset_mode = self.get_appliance.mode
         self._speed = self.get_entity.state
 
     @property
@@ -57,7 +59,10 @@ def speed_count(self) -> int:
     @property
     def percentage(self):
         """Return the current speed percentage."""
-        speed = self._speed if self.get_entity.state is None else self.get_entity.state
+        if self._preset_mode == WorkMode.OFF:
+            speed = 0
+        else:
+            speed = self._speed if self.get_entity.state is None else self.get_entity.state
         percentage = ranged_value_to_percentage(self._speed_range, speed)
         _LOGGER.debug(f"percentage - speed: {speed} percentage: {percentage}")
         return percentage
@@ -74,16 +79,16 @@ async def async_set_percentage(self, percentage: int) -> None:
             await self.async_turn_off()
             return
 
-        is_manual = self.preset_mode is Mode.MANUAL
+        is_manual = self.preset_mode == WorkMode.MANUAL
         # make sure manual is set before setting speed
         if not is_manual:
-            await self.async_set_preset_mode(Mode.MANUAL)
+            await self.async_set_preset_mode(WorkMode.MANUAL)
 
         await self.api.set_fan_speed(self.pnc_id, self._speed)
 
-        if is_manual:
-            await asyncio.sleep(10)
-            await self.coordinator.async_request_refresh()
+        self.async_write_ha_state()
+        await asyncio.sleep(10)
+        await self.coordinator.async_request_refresh()
 
     @property
     def supported_features(self):
@@ -93,29 +98,34 @@ def supported_features(self):
     @property
     def preset_mode(self):
         """Return the current preset mode, e.g., auto, smart, interval, favorite."""
-        return self._preset_mode if self.get_appliance.mode is Mode.UNDEFINED else self.get_appliance.mode
+        return (
+            self._preset_mode.value
+            if self.get_appliance.mode.value is WorkMode.UNDEFINED.value
+            else self.get_appliance.mode.value
+        )
 
     @property
     def preset_modes(self):
         """Return a list of available preset modes."""
-        return PRESET_MODES
+        return self.get_appliance.preset_modes
 
     async def async_set_preset_mode(self, preset_mode: str) -> None:
         """Set new preset mode."""
         self._valid_preset_mode_or_raise(preset_mode)
-        self._preset_mode = Mode(preset_mode)
-        self.get_appliance.clear_mode()
+        self._preset_mode = WorkMode(preset_mode)
+
+        self.get_appliance.set_mode(self._preset_mode)
         self.async_write_ha_state()
         await self.api.set_work_mode(self.pnc_id, self._preset_mode)
         await asyncio.sleep(10)
         await self.coordinator.async_request_refresh()
 
     @property
     def is_on(self):
-        return self.preset_mode is not Mode.OFF
+        return self.preset_mode is not WorkMode.OFF
 
     async def async_turn_on(self, percentage: int | None = None, preset_mode: str | None = None, **kwargs) -> None:
-        self._preset_mode = Mode(preset_mode or Mode.AUTO.value)
+        self._preset_mode = self.get_appliance.work_mode_from_preset_mode(preset_mode)
 
         # Handle incorrect percentage
         if percentage is not None and isinstance(percentage, str):
@@ -127,8 +137,7 @@ async def async_turn_on(self, percentage: int | None = None, preset_mode: str |
 
         # Proceed with the provided or default percentage
         self._speed = math.floor(percentage_to_ranged_value(self._speed_range, percentage or 10))
-        self.get_appliance.clear_mode()
-        self.get_entity.clear_state()
+        self.get_appliance.set_mode(self._preset_mode)
         self.async_write_ha_state()
 
         await self.api.set_work_mode(self.pnc_id, self._preset_mode)
@@ -138,12 +147,11 @@ async def async_turn_on(self, percentage: int | None = None, preset_mode: str |
 
     async def async_turn_off(self, **kwargs) -> None:
         """Turn off the entity."""
-        self._preset_mode = Mode.OFF
+        self._preset_mode = WorkMode.OFF
         self._speed = 0
-        self.get_appliance.clear_mode()
-        self.get_entity.clear_state()
+        self.get_appliance.set_mode(self._preset_mode)
         self.async_write_ha_state()
 
-        await self.api.set_work_mode(self.pnc_id, Mode.OFF)
+        await self.api.set_work_mode(self.pnc_id, WorkMode.OFF)
         await asyncio.sleep(10)
         await self.coordinator.async_request_refresh()