Add humidifier reproduce state (#36799)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/humidifier/reproduce_state.py

@@ -0,0 +1,96 @@
+"""Module that groups code required to handle state restore for component."""
+import asyncio
+import logging
+from typing import Any, Dict, Iterable, Optional
+
+from homeassistant.const import SERVICE_TURN_OFF, SERVICE_TURN_ON, STATE_OFF, STATE_ON
+from homeassistant.core import Context, State
+from homeassistant.helpers.typing import HomeAssistantType
+
+from .const import (
+    ATTR_HUMIDITY,
+    ATTR_MODE,
+    DOMAIN,
+    SERVICE_SET_HUMIDITY,
+    SERVICE_SET_MODE,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def _async_reproduce_states(
+    hass: HomeAssistantType,
+    state: State,
+    *,
+    context: Optional[Context] = None,
+    reproduce_options: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Reproduce component states."""
+    cur_state = hass.states.get(state.entity_id)
+
+    if cur_state is None:
+        _LOGGER.warning("Unable to find entity %s", state.entity_id)
+        return
+
+    async def call_service(service: str, keys: Iterable, data=None):
+        """Call service with set of attributes given."""
+        data = data or {}
+        data["entity_id"] = state.entity_id
+        for key in keys:
+            if key in state.attributes:
+                data[key] = state.attributes[key]
+
+        await hass.services.async_call(
+            DOMAIN, service, data, blocking=True, context=context
+        )
+
+    if state.state == STATE_OFF:
+        # Ensure the device is off if it needs to be and exit
+        if cur_state.state != STATE_OFF:
+            await call_service(SERVICE_TURN_OFF, [])
+        return
+
+    if state.state != STATE_ON:
+        # we can't know how to handle this
+        _LOGGER.warning(
+            "Invalid state specified for %s: %s", state.entity_id, state.state
+        )
+        return
+
+    # First of all, turn on if needed, because the device might not
+    # be able to set mode and humidity while being off
+    if cur_state.state != STATE_ON:
+        await call_service(SERVICE_TURN_ON, [])
+        # refetch the state as turning on might allow us to see some more values
+        cur_state = hass.states.get(state.entity_id)
+
+    # Then set the mode before target humidity, because switching modes
+    # may invalidate target humidity
+    if ATTR_MODE in state.attributes and state.attributes[
+        ATTR_MODE
+    ] != cur_state.attributes.get(ATTR_MODE):
+        await call_service(SERVICE_SET_MODE, [ATTR_MODE])
+
+    # Next, restore target humidity for the current mode
+    if ATTR_HUMIDITY in state.attributes and state.attributes[
+        ATTR_HUMIDITY
+    ] != cur_state.attributes.get(ATTR_HUMIDITY):
+        await call_service(SERVICE_SET_HUMIDITY, [ATTR_HUMIDITY])
+
+
+async def async_reproduce_states(
+    hass: HomeAssistantType,
+    states: Iterable[State],
+    *,
+    context: Optional[Context] = None,
+    reproduce_options: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Reproduce component states."""
+    await asyncio.gather(
+        *(
+            _async_reproduce_states(
+                hass, state, context=context, reproduce_options=reproduce_options
+            )
+            for state in states
+        )
+    )

tests/components/humidifier/test_reproduce_state.py

@@ -0,0 +1,237 @@
+"""The tests for reproduction of state."""
+
+import pytest
+
+from homeassistant.components.humidifier.const import (
+    ATTR_HUMIDITY,
+    ATTR_MODE,
+    DOMAIN,
+    MODE_AWAY,
+    MODE_ECO,
+    MODE_NORMAL,
+    SERVICE_SET_HUMIDITY,
+    SERVICE_SET_MODE,
+)
+from homeassistant.components.humidifier.reproduce_state import async_reproduce_states
+from homeassistant.const import SERVICE_TURN_OFF, SERVICE_TURN_ON, STATE_OFF, STATE_ON
+from homeassistant.core import Context, State
+
+from tests.common import async_mock_service
+
+ENTITY_1 = "humidifier.test1"
+ENTITY_2 = "humidifier.test2"
+
+
+async def test_reproducing_on_off_states(hass, caplog):
+    """Test reproducing humidifier states."""
+    hass.states.async_set(ENTITY_1, "off", {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 45})
+    hass.states.async_set(ENTITY_2, "on", {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 45})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    mode_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_MODE)
+    humidity_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_HUMIDITY)
+
+    # These calls should do nothing as entities already in desired state
+    await hass.helpers.state.async_reproduce_state(
+        [
+            State(ENTITY_1, "off", {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 45}),
+            State(ENTITY_2, "on", {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 45}),
+        ],
+    )
+
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 0
+    assert len(mode_calls) == 0
+    assert len(humidity_calls) == 0
+
+    # Test invalid state is handled
+    await hass.helpers.state.async_reproduce_state([State(ENTITY_1, "not_supported")])
+
+    assert "not_supported" in caplog.text
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 0
+    assert len(mode_calls) == 0
+    assert len(humidity_calls) == 0
+
+    # Make sure correct services are called
+    await hass.helpers.state.async_reproduce_state(
+        [
+            State(ENTITY_2, "off"),
+            State(ENTITY_1, "on", {}),
+            # Should not raise
+            State("humidifier.non_existing", "on"),
+        ]
+    )
+
+    assert len(turn_on_calls) == 1
+    assert turn_on_calls[0].domain == "humidifier"
+    assert turn_on_calls[0].data == {"entity_id": ENTITY_1}
+
+    assert len(turn_off_calls) == 1
+    assert turn_off_calls[0].domain == "humidifier"
+    assert turn_off_calls[0].data == {"entity_id": ENTITY_2}
+
+    # Make sure we didn't call services for missing attributes
+    assert len(mode_calls) == 0
+    assert len(humidity_calls) == 0
+
+
+async def test_multiple_attrs(hass):
+    """Test turn on with multiple attributes."""
+    hass.states.async_set(ENTITY_1, STATE_OFF, {})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    mode_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_MODE)
+    humidity_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_HUMIDITY)
+
+    await async_reproduce_states(
+        hass, [State(ENTITY_1, STATE_ON, {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 45})]
+    )
+
+    await hass.async_block_till_done()
+
+    assert len(turn_on_calls) == 1
+    assert turn_on_calls[0].data == {"entity_id": ENTITY_1}
+    assert len(turn_off_calls) == 0
+    assert len(mode_calls) == 1
+    assert mode_calls[0].data == {"entity_id": ENTITY_1, "mode": "normal"}
+    assert len(humidity_calls) == 1
+    assert humidity_calls[0].data == {"entity_id": ENTITY_1, "humidity": 45}
+
+
+async def test_turn_off_multiple_attrs(hass):
+    """Test set mode and humidity for off state."""
+    hass.states.async_set(ENTITY_1, STATE_ON, {})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    mode_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_MODE)
+    humidity_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_HUMIDITY)
+
+    await async_reproduce_states(
+        hass, [State(ENTITY_1, STATE_OFF, {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 45})]
+    )
+
+    await hass.async_block_till_done()
+
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 1
+    assert turn_off_calls[0].data == {"entity_id": ENTITY_1}
+    assert len(mode_calls) == 0
+    assert len(humidity_calls) == 0
+
+
+async def test_multiple_modes(hass):
+    """Test that multiple states gets calls."""
+    hass.states.async_set(ENTITY_1, STATE_OFF, {})
+    hass.states.async_set(ENTITY_2, STATE_OFF, {})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    mode_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_MODE)
+    humidity_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_HUMIDITY)
+
+    await async_reproduce_states(
+        hass,
+        [
+            State(ENTITY_1, STATE_ON, {ATTR_MODE: MODE_ECO, ATTR_HUMIDITY: 40}),
+            State(ENTITY_2, STATE_ON, {ATTR_MODE: MODE_NORMAL, ATTR_HUMIDITY: 50}),
+        ],
+    )
+
+    await hass.async_block_till_done()
+
+    assert len(turn_on_calls) == 2
+    assert len(turn_off_calls) == 0
+    assert len(mode_calls) == 2
+    # order is not guaranteed
+    assert any(
+        call.data == {"entity_id": ENTITY_1, "mode": MODE_ECO} for call in mode_calls
+    )
+    assert any(
+        call.data == {"entity_id": ENTITY_2, "mode": MODE_NORMAL} for call in mode_calls
+    )
+    assert len(humidity_calls) == 2
+    # order is not guaranteed
+    assert any(
+        call.data == {"entity_id": ENTITY_1, "humidity": 40} for call in humidity_calls
+    )
+    assert any(
+        call.data == {"entity_id": ENTITY_2, "humidity": 50} for call in humidity_calls
+    )
+
+
+async def test_state_with_none(hass):
+    """Test that none is not a humidifier state."""
+    hass.states.async_set(ENTITY_1, STATE_OFF, {})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    mode_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_MODE)
+    humidity_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_HUMIDITY)
+
+    await async_reproduce_states(hass, [State(ENTITY_1, None)])
+
+    await hass.async_block_till_done()
+
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 0
+    assert len(mode_calls) == 0
+    assert len(humidity_calls) == 0
+
+
+async def test_state_with_context(hass):
+    """Test that context is forwarded."""
+    hass.states.async_set(ENTITY_1, STATE_OFF, {})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    mode_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_MODE)
+    humidity_calls = async_mock_service(hass, DOMAIN, SERVICE_SET_HUMIDITY)
+
+    context = Context()
+
+    await async_reproduce_states(
+        hass,
+        [State(ENTITY_1, STATE_ON, {ATTR_MODE: MODE_AWAY, ATTR_HUMIDITY: 45})],
+        context=context,
+    )
+
+    await hass.async_block_till_done()
+
+    assert len(turn_on_calls) == 1
+    assert turn_on_calls[0].data == {"entity_id": ENTITY_1}
+    assert turn_on_calls[0].context == context
+    assert len(turn_off_calls) == 0
+    assert len(mode_calls) == 1
+    assert mode_calls[0].data == {"entity_id": ENTITY_1, "mode": "away"}
+    assert mode_calls[0].context == context
+    assert len(humidity_calls) == 1
+    assert humidity_calls[0].data == {"entity_id": ENTITY_1, "humidity": 45}
+    assert humidity_calls[0].context == context
+
+
+@pytest.mark.parametrize(
+    "service,attribute",
+    [(SERVICE_SET_MODE, ATTR_MODE), (SERVICE_SET_HUMIDITY, ATTR_HUMIDITY)],
+)
+async def test_attribute(hass, service, attribute):
+    """Test that service call is made for each attribute."""
+    hass.states.async_set(ENTITY_1, STATE_ON, {})
+
+    turn_on_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, DOMAIN, SERVICE_TURN_OFF)
+    calls_1 = async_mock_service(hass, DOMAIN, service)
+
+    value = "dummy"
+
+    await async_reproduce_states(hass, [State(ENTITY_1, STATE_ON, {attribute: value})])
+
+    await hass.async_block_till_done()
+
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 0
+    assert len(calls_1) == 1
+    assert calls_1[0].data == {"entity_id": ENTITY_1, attribute: value}