Vacuum reproduce state (#27868)

* Vacuum reproduce state

* Add missing states

* Improved process

* Fix tests

homeassistant/components/vacuum/reproduce_state.py

@@ -0,0 +1,101 @@
+"""Reproduce an Vacuum state."""
+import asyncio
+import logging
+from typing import Iterable, Optional
+
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_IDLE,
+    STATE_OFF,
+    STATE_ON,
+    STATE_PAUSED,
+)
+from homeassistant.core import Context, State
+from homeassistant.helpers.typing import HomeAssistantType
+
+from . import (
+    ATTR_FAN_SPEED,
+    DOMAIN,
+    SERVICE_PAUSE,
+    SERVICE_RETURN_TO_BASE,
+    SERVICE_SET_FAN_SPEED,
+    SERVICE_START,
+    SERVICE_STOP,
+    STATE_CLEANING,
+    STATE_DOCKED,
+    STATE_RETURNING,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+VALID_STATES_TOGGLE = {STATE_ON, STATE_OFF}
+VALID_STATES_STATE = {
+    STATE_CLEANING,
+    STATE_DOCKED,
+    STATE_IDLE,
+    STATE_RETURNING,
+    STATE_PAUSED,
+}
+
+
+async def _async_reproduce_state(
+    hass: HomeAssistantType, state: State, context: Optional[Context] = None
+) -> None:
+    """Reproduce a single state."""
+    cur_state = hass.states.get(state.entity_id)
+
+    if cur_state is None:
+        _LOGGER.warning("Unable to find entity %s", state.entity_id)
+        return
+
+    if state.state not in VALID_STATES_TOGGLE and state.state not in VALID_STATES_STATE:
+        _LOGGER.warning(
+            "Invalid state specified for %s: %s", state.entity_id, state.state
+        )
+        return
+
+    # Return if we are already at the right state.
+    if cur_state.state == state.state and cur_state.attributes.get(
+        ATTR_FAN_SPEED
+    ) == state.attributes.get(ATTR_FAN_SPEED):
+        return
+
+    service_data = {ATTR_ENTITY_ID: state.entity_id}
+
+    if cur_state.state != state.state:
+        # Wrong state
+        if state.state == STATE_ON:
+            service = SERVICE_TURN_ON
+        elif state.state == STATE_OFF:
+            service = SERVICE_TURN_OFF
+        elif state.state == STATE_CLEANING:
+            service = SERVICE_START
+        elif state.state == STATE_DOCKED or state.state == STATE_RETURNING:
+            service = SERVICE_RETURN_TO_BASE
+        elif state.state == STATE_IDLE:
+            service = SERVICE_STOP
+        elif state.state == STATE_PAUSED:
+            service = SERVICE_PAUSE
+
+        await hass.services.async_call(
+            DOMAIN, service, service_data, context=context, blocking=True
+        )
+
+    if cur_state.attributes.get(ATTR_FAN_SPEED) != state.attributes.get(ATTR_FAN_SPEED):
+        # Wrong fan speed
+        service_data["fan_speed"] = state.attributes[ATTR_FAN_SPEED]
+        await hass.services.async_call(
+            DOMAIN, SERVICE_SET_FAN_SPEED, service_data, context=context, blocking=True
+        )
+
+
+async def async_reproduce_states(
+    hass: HomeAssistantType, states: Iterable[State], context: Optional[Context] = None
+) -> None:
+    """Reproduce Vacuum states."""
+    # Reproduce states in parallel.
+    await asyncio.gather(
+        *(_async_reproduce_state(hass, state, context) for state in states)
+    )

tests/components/vacuum/test_reproduce_state.py

@@ -0,0 +1,139 @@
+"""Test reproduce state for Vacuum."""
+from homeassistant.components.vacuum import (
+    ATTR_FAN_SPEED,
+    SERVICE_PAUSE,
+    SERVICE_RETURN_TO_BASE,
+    SERVICE_SET_FAN_SPEED,
+    SERVICE_START,
+    SERVICE_STOP,
+    STATE_CLEANING,
+    STATE_DOCKED,
+    STATE_RETURNING,
+)
+from homeassistant.const import (
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_IDLE,
+    STATE_OFF,
+    STATE_ON,
+    STATE_PAUSED,
+)
+from homeassistant.core import State
+
+from tests.common import async_mock_service
+
+FAN_SPEED_LOW = "low"
+FAN_SPEED_HIGH = "high"
+
+
+async def test_reproducing_states(hass, caplog):
+    """Test reproducing Vacuum states."""
+    hass.states.async_set("vacuum.entity_off", STATE_OFF, {})
+    hass.states.async_set("vacuum.entity_on", STATE_ON, {})
+    hass.states.async_set(
+        "vacuum.entity_on_fan", STATE_ON, {ATTR_FAN_SPEED: FAN_SPEED_LOW}
+    )
+    hass.states.async_set("vacuum.entity_cleaning", STATE_CLEANING, {})
+    hass.states.async_set("vacuum.entity_docked", STATE_DOCKED, {})
+    hass.states.async_set("vacuum.entity_idle", STATE_IDLE, {})
+    hass.states.async_set("vacuum.entity_returning", STATE_RETURNING, {})
+    hass.states.async_set("vacuum.entity_paused", STATE_PAUSED, {})
+
+    turn_on_calls = async_mock_service(hass, "vacuum", SERVICE_TURN_ON)
+    turn_off_calls = async_mock_service(hass, "vacuum", SERVICE_TURN_OFF)
+    start_calls = async_mock_service(hass, "vacuum", SERVICE_START)
+    pause_calls = async_mock_service(hass, "vacuum", SERVICE_PAUSE)
+    stop_calls = async_mock_service(hass, "vacuum", SERVICE_STOP)
+    return_calls = async_mock_service(hass, "vacuum", SERVICE_RETURN_TO_BASE)
+    fan_speed_calls = async_mock_service(hass, "vacuum", SERVICE_SET_FAN_SPEED)
+
+    # These calls should do nothing as entities already in desired state
+    await hass.helpers.state.async_reproduce_state(
+        [
+            State("vacuum.entity_off", STATE_OFF),
+            State("vacuum.entity_on", STATE_ON),
+            State("vacuum.entity_on_fan", STATE_ON, {ATTR_FAN_SPEED: FAN_SPEED_LOW}),
+            State("vacuum.entity_cleaning", STATE_CLEANING),
+            State("vacuum.entity_docked", STATE_DOCKED),
+            State("vacuum.entity_idle", STATE_IDLE),
+            State("vacuum.entity_returning", STATE_RETURNING),
+            State("vacuum.entity_paused", STATE_PAUSED),
+        ],
+        blocking=True,
+    )
+
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 0
+    assert len(start_calls) == 0
+    assert len(pause_calls) == 0
+    assert len(stop_calls) == 0
+    assert len(return_calls) == 0
+    assert len(fan_speed_calls) == 0
+
+    # Test invalid state is handled
+    await hass.helpers.state.async_reproduce_state(
+        [State("vacuum.entity_off", "not_supported")], blocking=True
+    )
+
+    assert "not_supported" in caplog.text
+    assert len(turn_on_calls) == 0
+    assert len(turn_off_calls) == 0
+    assert len(start_calls) == 0
+    assert len(pause_calls) == 0
+    assert len(stop_calls) == 0
+    assert len(return_calls) == 0
+    assert len(fan_speed_calls) == 0
+
+    # Make sure correct services are called
+    await hass.helpers.state.async_reproduce_state(
+        [
+            State("vacuum.entity_off", STATE_ON),
+            State("vacuum.entity_on", STATE_OFF),
+            State("vacuum.entity_on_fan", STATE_ON, {ATTR_FAN_SPEED: FAN_SPEED_HIGH}),
+            State("vacuum.entity_cleaning", STATE_PAUSED),
+            State("vacuum.entity_docked", STATE_CLEANING),
+            State("vacuum.entity_idle", STATE_DOCKED),
+            State("vacuum.entity_returning", STATE_CLEANING),
+            State("vacuum.entity_paused", STATE_IDLE),
+            # Should not raise
+            State("vacuum.non_existing", STATE_ON),
+        ],
+        blocking=True,
+    )
+
+    assert len(turn_on_calls) == 1
+    assert turn_on_calls[0].domain == "vacuum"
+    assert turn_on_calls[0].data == {"entity_id": "vacuum.entity_off"}
+
+    assert len(turn_off_calls) == 1
+    assert turn_off_calls[0].domain == "vacuum"
+    assert turn_off_calls[0].data == {"entity_id": "vacuum.entity_on"}
+
+    assert len(start_calls) == 2
+    entities = [
+        {"entity_id": "vacuum.entity_docked"},
+        {"entity_id": "vacuum.entity_returning"},
+    ]
+    for call in start_calls:
+        assert call.domain == "vacuum"
+        assert call.data in entities
+        entities.remove(call.data)
+
+    assert len(pause_calls) == 1
+    assert pause_calls[0].domain == "vacuum"
+    assert pause_calls[0].data == {"entity_id": "vacuum.entity_cleaning"}
+
+    assert len(stop_calls) == 1
+    assert stop_calls[0].domain == "vacuum"
+    assert stop_calls[0].data == {"entity_id": "vacuum.entity_paused"}
+
+    assert len(return_calls) == 1
+    assert return_calls[0].domain == "vacuum"
+    assert return_calls[0].data == {"entity_id": "vacuum.entity_idle"}
+
+    assert len(fan_speed_calls) == 1
+    assert fan_speed_calls[0].domain == "vacuum"
+    assert fan_speed_calls[0].data == {
+        "entity_id": "vacuum.entity_on_fan",
+        ATTR_FAN_SPEED: FAN_SPEED_HIGH,
+    }