diff --git a/homeassistant/components/vacuum/reproduce_state.py b/homeassistant/components/vacuum/reproduce_state.py new file mode 100644 index 00000000000..485ffef0c9f --- /dev/null +++ b/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) + ) diff --git a/tests/components/vacuum/test_reproduce_state.py b/tests/components/vacuum/test_reproduce_state.py new file mode 100644 index 00000000000..d5a7051e6a6 --- /dev/null +++ b/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, + }