diff --git a/homeassistant/components/template/vacuum.py b/homeassistant/components/template/vacuum.py index 5ff99020f0d..67f0f780388 100644 --- a/homeassistant/components/template/vacuum.py +++ b/homeassistant/components/template/vacuum.py @@ -44,6 +44,7 @@ from .helpers import async_setup_template_platform from .template_entity import ( TEMPLATE_ENTITY_ATTRIBUTES_SCHEMA_LEGACY, TEMPLATE_ENTITY_AVAILABILITY_SCHEMA_LEGACY, + TEMPLATE_ENTITY_OPTIMISTIC_SCHEMA, TemplateEntity, make_template_entity_common_modern_attributes_schema, ) @@ -76,24 +77,26 @@ LEGACY_FIELDS = { CONF_VALUE_TEMPLATE: CONF_STATE, } -VACUUM_YAML_SCHEMA = vol.All( - vol.Schema( - { - vol.Optional(CONF_BATTERY_LEVEL): cv.template, - vol.Optional(CONF_FAN_SPEED_LIST, default=[]): cv.ensure_list, - vol.Optional(CONF_FAN_SPEED): cv.template, - vol.Optional(CONF_STATE): cv.template, - vol.Optional(SERVICE_CLEAN_SPOT): cv.SCRIPT_SCHEMA, - vol.Optional(SERVICE_LOCATE): cv.SCRIPT_SCHEMA, - vol.Optional(SERVICE_PAUSE): cv.SCRIPT_SCHEMA, - vol.Optional(SERVICE_RETURN_TO_BASE): cv.SCRIPT_SCHEMA, - vol.Optional(SERVICE_SET_FAN_SPEED): cv.SCRIPT_SCHEMA, - vol.Required(SERVICE_START): cv.SCRIPT_SCHEMA, - vol.Optional(SERVICE_STOP): cv.SCRIPT_SCHEMA, - } - ).extend(make_template_entity_common_modern_attributes_schema(DEFAULT_NAME).schema) +VACUUM_COMMON_SCHEMA = vol.Schema( + { + vol.Optional(CONF_BATTERY_LEVEL): cv.template, + vol.Optional(CONF_FAN_SPEED_LIST, default=[]): cv.ensure_list, + vol.Optional(CONF_FAN_SPEED): cv.template, + vol.Optional(CONF_STATE): cv.template, + vol.Optional(SERVICE_CLEAN_SPOT): cv.SCRIPT_SCHEMA, + vol.Optional(SERVICE_LOCATE): cv.SCRIPT_SCHEMA, + vol.Optional(SERVICE_PAUSE): cv.SCRIPT_SCHEMA, + vol.Optional(SERVICE_RETURN_TO_BASE): cv.SCRIPT_SCHEMA, + vol.Optional(SERVICE_SET_FAN_SPEED): cv.SCRIPT_SCHEMA, + vol.Required(SERVICE_START): cv.SCRIPT_SCHEMA, + vol.Optional(SERVICE_STOP): cv.SCRIPT_SCHEMA, + } ) +VACUUM_YAML_SCHEMA = VACUUM_COMMON_SCHEMA.extend( + TEMPLATE_ENTITY_OPTIMISTIC_SCHEMA +).extend(make_template_entity_common_modern_attributes_schema(DEFAULT_NAME).schema) + VACUUM_LEGACY_YAML_SCHEMA = vol.All( cv.deprecated(CONF_ENTITY_ID), vol.Schema( @@ -147,16 +150,15 @@ class AbstractTemplateVacuum(AbstractTemplateEntity, StateVacuumEntity): """Representation of a template vacuum features.""" _entity_id_format = ENTITY_ID_FORMAT + _optimistic_entity = True # The super init is not called because TemplateEntity and TriggerEntity will call AbstractTemplateEntity.__init__. # This ensures that the __init__ on AbstractTemplateEntity is not called twice. def __init__(self, config: dict[str, Any]) -> None: # pylint: disable=super-init-not-called """Initialize the features.""" - self._template = config.get(CONF_STATE) self._battery_level_template = config.get(CONF_BATTERY_LEVEL) self._fan_speed_template = config.get(CONF_FAN_SPEED) - self._state = None self._battery_level = None self._attr_fan_speed = None @@ -185,17 +187,12 @@ class AbstractTemplateVacuum(AbstractTemplateEntity, StateVacuumEntity): if (action_config := config.get(action_id)) is not None: yield (action_id, action_config, supported_feature) - @property - def activity(self) -> VacuumActivity | None: - """Return the status of the vacuum cleaner.""" - return self._state - def _handle_state(self, result: Any) -> None: # Validate state if result in _VALID_STATES: - self._state = result + self._attr_activity = result elif result == STATE_UNKNOWN: - self._state = None + self._attr_activity = None else: _LOGGER.error( "Received invalid vacuum state: %s for entity %s. Expected: %s", @@ -203,31 +200,46 @@ class AbstractTemplateVacuum(AbstractTemplateEntity, StateVacuumEntity): self.entity_id, ", ".join(_VALID_STATES), ) - self._state = None + self._attr_activity = None async def async_start(self) -> None: """Start or resume the cleaning task.""" + if self._attr_assumed_state: + self._attr_activity = VacuumActivity.CLEANING + self.async_write_ha_state() await self.async_run_script( self._action_scripts[SERVICE_START], context=self._context ) async def async_pause(self) -> None: """Pause the cleaning task.""" + if self._attr_assumed_state: + self._attr_activity = VacuumActivity.PAUSED + self.async_write_ha_state() if script := self._action_scripts.get(SERVICE_PAUSE): await self.async_run_script(script, context=self._context) async def async_stop(self, **kwargs: Any) -> None: """Stop the cleaning task.""" + if self._attr_assumed_state: + self._attr_activity = VacuumActivity.IDLE + self.async_write_ha_state() if script := self._action_scripts.get(SERVICE_STOP): await self.async_run_script(script, context=self._context) async def async_return_to_base(self, **kwargs: Any) -> None: """Set the vacuum cleaner to return to the dock.""" + if self._attr_assumed_state: + self._attr_activity = VacuumActivity.RETURNING + self.async_write_ha_state() if script := self._action_scripts.get(SERVICE_RETURN_TO_BASE): await self.async_run_script(script, context=self._context) async def async_clean_spot(self, **kwargs: Any) -> None: """Perform a spot clean-up.""" + if self._attr_assumed_state: + self._attr_activity = VacuumActivity.CLEANING + self.async_write_ha_state() if script := self._action_scripts.get(SERVICE_CLEAN_SPOT): await self.async_run_script(script, context=self._context) @@ -274,7 +286,7 @@ class AbstractTemplateVacuum(AbstractTemplateEntity, StateVacuumEntity): if isinstance(fan_speed, TemplateError): # This is legacy behavior self._attr_fan_speed = None - self._state = None + self._attr_activity = None return if fan_speed in self._attr_fan_speed_list: @@ -320,7 +332,7 @@ class TemplateStateVacuumEntity(TemplateEntity, AbstractTemplateVacuum): """Set up templates.""" if self._template is not None: self.add_template_attribute( - "_state", self._template, None, self._update_state + "_attr_activity", self._template, None, self._update_state ) if self._fan_speed_template is not None: self.add_template_attribute( @@ -344,7 +356,7 @@ class TemplateStateVacuumEntity(TemplateEntity, AbstractTemplateVacuum): super()._update_state(result) if isinstance(result, TemplateError): # This is legacy behavior - self._state = STATE_UNKNOWN + self._attr_activity = None if not self._availability_template: self._attr_available = True return diff --git a/tests/components/template/test_vacuum.py b/tests/components/template/test_vacuum.py index ae65823309a..540b4eccd3b 100644 --- a/tests/components/template/test_vacuum.py +++ b/tests/components/template/test_vacuum.py @@ -1153,3 +1153,111 @@ async def test_empty_action_config( assert state.attributes["supported_features"] == ( VacuumEntityFeature.STATE | VacuumEntityFeature.START | supported_features ) + + +@pytest.mark.parametrize( + ("count", "vacuum_config"), + [ + ( + 1, + {"name": TEST_OBJECT_ID, "start": [], **TEMPLATE_VACUUM_ACTIONS}, + ) + ], +) +@pytest.mark.parametrize( + "style", + [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER], +) +@pytest.mark.parametrize( + ("service", "expected"), + [ + (vacuum.SERVICE_START, VacuumActivity.CLEANING), + (vacuum.SERVICE_PAUSE, VacuumActivity.PAUSED), + (vacuum.SERVICE_STOP, VacuumActivity.IDLE), + (vacuum.SERVICE_RETURN_TO_BASE, VacuumActivity.RETURNING), + (vacuum.SERVICE_CLEAN_SPOT, VacuumActivity.CLEANING), + ], +) +@pytest.mark.usefixtures("setup_vacuum") +async def test_assumed_optimistic( + hass: HomeAssistant, + service: str, + expected: VacuumActivity, + calls: list[ServiceCall], +) -> None: + """Test assumed optimistic.""" + + await hass.services.async_call( + vacuum.DOMAIN, + service, + {"entity_id": TEST_ENTITY_ID}, + blocking=True, + ) + await hass.async_block_till_done() + + state = hass.states.get(TEST_ENTITY_ID) + assert state.state == expected + + +@pytest.mark.parametrize( + ("count", "vacuum_config"), + [ + ( + 1, + { + "name": TEST_OBJECT_ID, + "state": "{{ states('sensor.test_state') }}", + "start": [], + **TEMPLATE_VACUUM_ACTIONS, + "optimistic": True, + }, + ) + ], +) +@pytest.mark.parametrize( + "style", + [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER], +) +@pytest.mark.parametrize( + ("service", "expected"), + [ + (vacuum.SERVICE_START, VacuumActivity.CLEANING), + (vacuum.SERVICE_PAUSE, VacuumActivity.PAUSED), + (vacuum.SERVICE_STOP, VacuumActivity.IDLE), + (vacuum.SERVICE_RETURN_TO_BASE, VacuumActivity.RETURNING), + (vacuum.SERVICE_CLEAN_SPOT, VacuumActivity.CLEANING), + ], +) +@pytest.mark.usefixtures("setup_vacuum") +async def test_optimistic_option( + hass: HomeAssistant, + service: str, + expected: VacuumActivity, + calls: list[ServiceCall], +) -> None: + """Test optimistic yaml option.""" + hass.states.async_set(TEST_STATE_SENSOR, VacuumActivity.DOCKED) + await hass.async_block_till_done() + + state = hass.states.get(TEST_ENTITY_ID) + assert state.state == VacuumActivity.DOCKED + + await hass.services.async_call( + vacuum.DOMAIN, + service, + {"entity_id": TEST_ENTITY_ID}, + blocking=True, + ) + await hass.async_block_till_done() + + state = hass.states.get(TEST_ENTITY_ID) + assert state.state == expected + + hass.states.async_set(TEST_STATE_SENSOR, VacuumActivity.RETURNING) + await hass.async_block_till_done() + + hass.states.async_set(TEST_STATE_SENSOR, VacuumActivity.DOCKED) + await hass.async_block_till_done() + + state = hass.states.get(TEST_ENTITY_ID) + assert state.state == VacuumActivity.DOCKED