Add control for preset modes of climate device (close #450)

dext0r · Nov 10, 2024 · ffe4914 · ffe4914
1 parent 4a62afe
commit ffe4914
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 0 deletions.
diff --git a/custom_components/yandex_smart_home/capability_mode.py b/custom_components/yandex_smart_home/capability_mode.py
@@ -381,6 +381,64 @@ class ProgramCapability(StateModeCapability, ABC):
     instance = ModeCapabilityInstance.PROGRAM
 
 
+class ProgramCapabilityClimate(ProgramCapability):
+    """Capability to control the mode preset of a climate device."""
+
+    _modes_map_default = {
+        ModeCapabilityMode.AUTO: [
+            climate.const.PRESET_NONE,
+        ],
+        ModeCapabilityMode.ECO: [
+            climate.const.PRESET_ECO,
+        ],
+        ModeCapabilityMode.MIN: [
+            climate.const.PRESET_AWAY,
+        ],
+        ModeCapabilityMode.TURBO: [
+            climate.const.PRESET_BOOST,
+        ],
+        ModeCapabilityMode.MEDIUM: [
+            climate.const.PRESET_COMFORT,
+        ],
+        ModeCapabilityMode.MAX: [
+            climate.const.PRESET_HOME,
+        ],
+        ModeCapabilityMode.QUIET: [
+            climate.const.PRESET_SLEEP,
+        ],
+    }
+
+    @property
+    def supported(self) -> bool:
+        """Test if the capability is supported."""
+        if self.state.domain == climate.DOMAIN and self._state_features & ClimateEntityFeature.PRESET_MODE:
+            return super().supported
+        return False
+
+    async def set_instance_state(self, context: Context, state: ModeCapabilityInstanceActionState) -> None:
+        """Change the capability state."""
+        await self._hass.services.async_call(
+            climate.DOMAIN,
+            climate.SERVICE_SET_PRESET_MODE,
+            {
+                ATTR_ENTITY_ID: self.state.entity_id,
+                climate.ATTR_PRESET_MODE: self.get_ha_mode_by_yandex_mode(state.value),
+            },
+            blocking=True,
+            context=context,
+        )
+
+    @property
+    def _ha_modes(self) -> Iterable[Any]:
+        """Returns list of HA modes."""
+        return self.state.attributes.get(climate.ATTR_PRESET_MODES, []) or []
+
+    @property
+    def _ha_value(self) -> Any:
+        """Return the current unmapped capability value."""
+        return self.state.attributes.get(climate.ATTR_PRESET_MODE)
+
+
 class ProgramCapabilityHumidifier(ProgramCapability):
     """Capability to control the mode of a humidifier device."""
 
@@ -903,6 +961,7 @@ def _ha_value(self) -> Any:
 
 STATE_CAPABILITIES_REGISTRY.register(ThermostatCapability)
 STATE_CAPABILITIES_REGISTRY.register(SwingCapability)
+STATE_CAPABILITIES_REGISTRY.register(ProgramCapabilityClimate)
 STATE_CAPABILITIES_REGISTRY.register(ProgramCapabilityHumidifier)
 STATE_CAPABILITIES_REGISTRY.register(ProgramCapabilityFan)
 STATE_CAPABILITIES_REGISTRY.register(InputSourceCapability)

diff --git a/tests/test_capability.py b/tests/test_capability.py
@@ -116,6 +116,7 @@ async def test_capability_demo_platform(hass: HomeAssistant, entry_data: MockCon
     assert capabilities == [
         ("devices.capabilities.mode", "thermostat"),
         ("devices.capabilities.mode", "swing"),
+        ("devices.capabilities.mode", "program"),
         ("devices.capabilities.mode", "fan_speed"),
         ("devices.capabilities.on_off", "on"),
     ]

diff --git a/tests/test_capability_mode.py b/tests/test_capability_mode.py
@@ -301,6 +301,50 @@ async def test_capability_mode_swing(hass: HomeAssistant, entry_data: MockConfig
     assert calls[0].data == {ATTR_ENTITY_ID: state.entity_id, climate.ATTR_SWING_MODE: "lr"}
 
 
+async def test_capability_mode_program_climate(hass: HomeAssistant, entry_data: MockConfigEntryData) -> None:
+    state = State("climate.test", STATE_OFF)
+    assert_no_capabilities(hass, entry_data, state, CapabilityType.MODE, ModeCapabilityInstance.PROGRAM)
+    state = State(
+        "climate.test",
+        STATE_OFF,
+        {
+            ATTR_SUPPORTED_FEATURES: climate.ClimateEntityFeature.PRESET_MODE,
+            climate.ATTR_PRESET_MODES: ["none", "eco", "boost"],
+        },
+    )
+    cap = cast(
+        ModeCapability,
+        get_exact_one_capability(hass, entry_data, state, CapabilityType.MODE, ModeCapabilityInstance.PROGRAM),
+    )
+    assert cap.retrievable is True
+    assert cap.parameters.dict() == {
+        "instance": "program",
+        "modes": [{"value": "auto"}, {"value": "eco"}, {"value": "turbo"}],
+    }
+    assert cap.get_value() is None
+    state = State(
+        "climate.test",
+        STATE_OFF,
+        {
+            ATTR_SUPPORTED_FEATURES: climate.ClimateEntityFeature.PRESET_MODE,
+            climate.ATTR_PRESET_MODES: ["none", "eco", "boost"],
+            climate.ATTR_PRESET_MODE: "eco",
+        },
+    )
+    cap = cast(
+        ModeCapability,
+        get_exact_one_capability(hass, entry_data, state, CapabilityType.MODE, ModeCapabilityInstance.PROGRAM),
+    )
+    assert cap.get_value() == ModeCapabilityMode.ECO
+    calls = async_mock_service(hass, climate.DOMAIN, climate.SERVICE_SET_PRESET_MODE)
+    await cap.set_instance_state(
+        Context(),
+        ModeCapabilityInstanceActionState(instance=ModeCapabilityInstance.PROGRAM, value=ModeCapabilityMode.TURBO),
+    )
+    assert len(calls) == 1
+    assert calls[0].data == {ATTR_ENTITY_ID: state.entity_id, climate.ATTR_PRESET_MODE: "boost"}
+
+
 async def test_capability_mode_program_humidifier(hass: HomeAssistant, entry_data: MockConfigEntryData) -> None:
     state = State("humidifier.test", STATE_OFF)
     assert_no_capabilities(hass, entry_data, state, CapabilityType.MODE, ModeCapabilityInstance.PROGRAM)