Add assumed optimistic to template number entities (#148499)

homeassistant/components/template/number.py

@@ -17,14 +17,9 @@ from homeassistant.components.number import (
     NumberEntity,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
+from homeassistant.const import CONF_NAME, CONF_STATE, CONF_UNIT_OF_MEASUREMENT
-    CONF_NAME,
-    CONF_OPTIMISTIC,
-    CONF_STATE,
-    CONF_UNIT_OF_MEASUREMENT,
-)
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers import config_validation as cv, template
 from homeassistant.helpers.entity_platform import (
     AddConfigEntryEntitiesCallback,
     AddEntitiesCallback,
@@ -33,6 +28,7 @@ from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from . import TriggerUpdateCoordinator
 from .const import CONF_MAX, CONF_MIN, CONF_STEP, DOMAIN
+from .entity import AbstractTemplateEntity
 from .helpers import (
     async_setup_template_entry,
     async_setup_template_platform,
@@ -40,6 +36,7 @@ from .helpers import (
 )
 from .template_entity import (
     TEMPLATE_ENTITY_COMMON_CONFIG_ENTRY_SCHEMA,
+    TEMPLATE_ENTITY_OPTIMISTIC_SCHEMA,
     TemplateEntity,
     make_template_entity_common_modern_schema,
 )
@@ -57,21 +54,15 @@ NUMBER_COMMON_SCHEMA = vol.Schema(
         vol.Optional(CONF_MAX, default=DEFAULT_MAX_VALUE): cv.template,
         vol.Optional(CONF_MIN, default=DEFAULT_MIN_VALUE): cv.template,
         vol.Required(CONF_SET_VALUE): cv.SCRIPT_SCHEMA,
-        vol.Required(CONF_STATE): cv.template,
+        vol.Optional(CONF_STATE): cv.template,
-        vol.Required(CONF_STEP): cv.template,
+        vol.Optional(CONF_STEP, default=DEFAULT_STEP): cv.template,
         vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
     }
-)
+).extend(make_template_entity_common_modern_schema(DEFAULT_NAME).schema)
 
-NUMBER_YAML_SCHEMA = (
+NUMBER_YAML_SCHEMA = NUMBER_COMMON_SCHEMA.extend(
-    vol.Schema(
+    TEMPLATE_ENTITY_OPTIMISTIC_SCHEMA
-        {
+).extend(make_template_entity_common_modern_schema(DEFAULT_NAME).schema)
-            vol.Optional(CONF_OPTIMISTIC, default=DEFAULT_OPTIMISTIC): cv.boolean,
-        }
-    )
-    .extend(make_template_entity_common_modern_schema(DEFAULT_NAME).schema)
-    .extend(NUMBER_COMMON_SCHEMA.schema)
-)
 
 NUMBER_CONFIG_ENTRY_SCHEMA = NUMBER_COMMON_SCHEMA.extend(
     TEMPLATE_ENTITY_COMMON_CONFIG_ENTRY_SCHEMA.schema
@@ -121,69 +112,28 @@ def async_create_preview_number(
     )
 
 
-class StateNumberEntity(TemplateEntity, NumberEntity):
+class AbstractTemplateNumber(AbstractTemplateEntity, NumberEntity):
-    """Representation of a template number."""
+    """Representation of a template number features."""
 
-    _attr_should_poll = False
     _entity_id_format = ENTITY_ID_FORMAT
+    _optimistic_entity = True
 
-    def __init__(
+    # The super init is not called because TemplateEntity and TriggerEntity will call AbstractTemplateEntity.__init__.
-        self,
+    # This ensures that the __init__ on AbstractTemplateEntity is not called twice.
-        hass: HomeAssistant,
+    def __init__(self, config: dict[str, Any]) -> None:  # pylint: disable=super-init-not-called
-        config,
+        """Initialize the features."""
-        unique_id: str | None,
-    ) -> None:
-        """Initialize the number."""
-        TemplateEntity.__init__(self, hass, config, unique_id)
-        if TYPE_CHECKING:
-            assert self._attr_name is not None
-
-        self._value_template = config[CONF_STATE]
-        self.add_script(CONF_SET_VALUE, config[CONF_SET_VALUE], self._attr_name, DOMAIN)
-
         self._step_template = config[CONF_STEP]
-        self._min_value_template = config[CONF_MIN]
+        self._min_template = config[CONF_MIN]
-        self._max_value_template = config[CONF_MAX]
+        self._max_template = config[CONF_MAX]
-        self._attr_assumed_state = self._optimistic = config.get(CONF_OPTIMISTIC)
+
         self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
         self._attr_native_step = DEFAULT_STEP
         self._attr_native_min_value = DEFAULT_MIN_VALUE
         self._attr_native_max_value = DEFAULT_MAX_VALUE
 
-    @callback
-    def _async_setup_templates(self) -> None:
-        """Set up templates."""
-        self.add_template_attribute(
-            "_attr_native_value",
-            self._value_template,
-            validator=vol.Coerce(float),
-            none_on_template_error=True,
-        )
-        self.add_template_attribute(
-            "_attr_native_step",
-            self._step_template,
-            validator=vol.Coerce(float),
-            none_on_template_error=True,
-        )
-        if self._min_value_template is not None:
-            self.add_template_attribute(
-                "_attr_native_min_value",
-                self._min_value_template,
-                validator=vol.Coerce(float),
-                none_on_template_error=True,
-            )
-        if self._max_value_template is not None:
-            self.add_template_attribute(
-                "_attr_native_max_value",
-                self._max_value_template,
-                validator=vol.Coerce(float),
-                none_on_template_error=True,
-            )
-        super()._async_setup_templates()
-
     async def async_set_native_value(self, value: float) -> None:
         """Set value of the number."""
-        if self._optimistic:
+        if self._attr_assumed_state:
             self._attr_native_value = value
             self.async_write_ha_state()
         if set_value := self._action_scripts.get(CONF_SET_VALUE):
@@ -194,17 +144,65 @@ class StateNumberEntity(TemplateEntity, NumberEntity):
             )
 
 
-class TriggerNumberEntity(TriggerEntity, NumberEntity):
+class StateNumberEntity(TemplateEntity, AbstractTemplateNumber):
+    """Representation of a template number."""
+
+    _attr_should_poll = False
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config: ConfigType,
+        unique_id: str | None,
+    ) -> None:
+        """Initialize the number."""
+        TemplateEntity.__init__(self, hass, config, unique_id)
+        AbstractTemplateNumber.__init__(self, config)
+
+        name = self._attr_name
+        if TYPE_CHECKING:
+            assert name is not None
+
+        self.add_script(CONF_SET_VALUE, config[CONF_SET_VALUE], name, DOMAIN)
+
+    @callback
+    def _async_setup_templates(self) -> None:
+        """Set up templates."""
+        if self._template is not None:
+            self.add_template_attribute(
+                "_attr_native_value",
+                self._template,
+                vol.Coerce(float),
+                none_on_template_error=True,
+            )
+        if self._step_template is not None:
+            self.add_template_attribute(
+                "_attr_native_step",
+                self._step_template,
+                vol.Coerce(float),
+                none_on_template_error=True,
+            )
+        if self._min_template is not None:
+            self.add_template_attribute(
+                "_attr_native_min_value",
+                self._min_template,
+                validator=vol.Coerce(float),
+                none_on_template_error=True,
+            )
+        if self._max_template is not None:
+            self.add_template_attribute(
+                "_attr_native_max_value",
+                self._max_template,
+                validator=vol.Coerce(float),
+                none_on_template_error=True,
+            )
+        super()._async_setup_templates()
+
+
+class TriggerNumberEntity(TriggerEntity, AbstractTemplateNumber):
     """Number entity based on trigger data."""
 
-    _entity_id_format = ENTITY_ID_FORMAT
     domain = NUMBER_DOMAIN
-    extra_template_keys = (
-        CONF_STATE,
-        CONF_STEP,
-        CONF_MIN,
-        CONF_MAX,
-    )
 
     def __init__(
         self,
@@ -213,47 +211,49 @@ class TriggerNumberEntity(TriggerEntity, NumberEntity):
         config: dict,
     ) -> None:
         """Initialize the entity."""
-        super().__init__(hass, coordinator, config)
+        TriggerEntity.__init__(self, hass, coordinator, config)
+        AbstractTemplateNumber.__init__(self, config)
 
-        name = self._rendered.get(CONF_NAME, DEFAULT_NAME)
+        for key in (
-        self.add_script(CONF_SET_VALUE, config[CONF_SET_VALUE], name, DOMAIN)
+            CONF_STATE,
+            CONF_STEP,
+            CONF_MIN,
+            CONF_MAX,
+        ):
+            if isinstance(config.get(key), template.Template):
+                self._to_render_simple.append(key)
+                self._parse_result.add(key)
 
-        self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
+        self.add_script(
-
+            CONF_SET_VALUE,
-    @property
+            config[CONF_SET_VALUE],
-    def native_value(self) -> float | None:
+            self._rendered.get(CONF_NAME, DEFAULT_NAME),
-        """Return the currently selected option."""
+            DOMAIN,
-        return vol.Any(vol.Coerce(float), None)(self._rendered.get(CONF_STATE))
-
-    @property
-    def native_min_value(self) -> int:
-        """Return the minimum value."""
-        return vol.Any(vol.Coerce(float), None)(
-            self._rendered.get(CONF_MIN, super().native_min_value)
         )
 
-    @property
+    def _handle_coordinator_update(self):
-    def native_max_value(self) -> int:
+        """Handle updated data from the coordinator."""
-        """Return the maximum value."""
+        self._process_data()
-        return vol.Any(vol.Coerce(float), None)(
-            self._rendered.get(CONF_MAX, super().native_max_value)
-        )
 
-    @property
+        if not self.available:
-    def native_step(self) -> int:
+            self.async_write_ha_state()
-        """Return the increment/decrement step."""
+            return
-        return vol.Any(vol.Coerce(float), None)(
+
-            self._rendered.get(CONF_STEP, super().native_step)
+        write_ha_state = False
-        )
+        for key, attr in (
-
+            (CONF_STATE, "_attr_native_value"),
-    async def async_set_native_value(self, value: float) -> None:
+            (CONF_STEP, "_attr_native_step"),
-        """Set value of the number."""
+            (CONF_MIN, "_attr_native_min_value"),
-        if self._config[CONF_OPTIMISTIC]:
+            (CONF_MAX, "_attr_native_max_value"),
-            self._attr_native_value = value
+        ):
+            if (rendered := self._rendered.get(key)) is not None:
+                setattr(self, attr, vol.Any(vol.Coerce(float), None)(rendered))
+                write_ha_state = True
+
+        if len(self._rendered) > 0:
+            # In case any non optimistic template
+            write_ha_state = True
+
+        if write_ha_state:
+            self.async_set_context(self.coordinator.data["context"])
             self.async_write_ha_state()
-        if set_value := self._action_scripts.get(CONF_SET_VALUE):
-            await self.async_run_script(
-                set_value,
-                run_variables={ATTR_VALUE: value},
-                context=self._context,
-            )

tests/components/template/test_number.py

@@ -29,6 +29,7 @@ from homeassistant.const import (
     CONF_ENTITY_ID,
     CONF_ICON,
     CONF_UNIT_OF_MEASUREMENT,
+    STATE_UNAVAILABLE,
     STATE_UNKNOWN,
 )
 from homeassistant.core import Context, HomeAssistant, ServiceCall
@@ -63,11 +64,11 @@ _VALUE_INPUT_NUMBER_CONFIG = {
 }
 
 TEST_STATE_ENTITY_ID = "number.test_state"
-
+TEST_AVAILABILITY_ENTITY_ID = "binary_sensor.test_availability"
 TEST_STATE_TRIGGER = {
     "trigger": {
         "trigger": "state",
-        "entity_id": [TEST_STATE_ENTITY_ID],
+        "entity_id": [TEST_STATE_ENTITY_ID, TEST_AVAILABILITY_ENTITY_ID],
     },
     "variables": {"triggering_entity": "{{ trigger.entity_id }}"},
     "action": [
@@ -191,19 +192,6 @@ async def test_missing_optional_config(hass: HomeAssistant) -> None:
 
 async def test_missing_required_keys(hass: HomeAssistant) -> None:
     """Test: missing required fields will fail."""
-    with assert_setup_component(0, "template"):
-        assert await setup.async_setup_component(
-            hass,
-            "template",
-            {
-                "template": {
-                    "number": {
-                        "set_value": {"service": "script.set_value"},
-                    }
-                }
-            },
-        )
-
     with assert_setup_component(0, "template"):
         assert await setup.async_setup_component(
             hass,
@@ -578,6 +566,91 @@ async def test_device_id(
     assert template_entity.device_id == device_entry.id
 
 
+@pytest.mark.parametrize(
+    ("count", "number_config"),
+    [
+        (
+            1,
+            {
+                "set_value": [],
+            },
+        )
+    ],
+)
+@pytest.mark.parametrize(
+    "style",
+    [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER],
+)
+@pytest.mark.usefixtures("setup_number")
+async def test_optimistic(hass: HomeAssistant) -> None:
+    """Test configuration with optimistic state."""
+    await hass.services.async_call(
+        number.DOMAIN,
+        number.SERVICE_SET_VALUE,
+        {ATTR_ENTITY_ID: _TEST_NUMBER, "value": 4},
+        blocking=True,
+    )
+
+    state = hass.states.get(_TEST_NUMBER)
+    assert float(state.state) == 4
+
+    await hass.services.async_call(
+        number.DOMAIN,
+        number.SERVICE_SET_VALUE,
+        {ATTR_ENTITY_ID: _TEST_NUMBER, "value": 2},
+        blocking=True,
+    )
+
+    state = hass.states.get(_TEST_NUMBER)
+    assert float(state.state) == 2
+
+
+@pytest.mark.parametrize(
+    ("count", "number_config"),
+    [
+        (
+            1,
+            {
+                "set_value": [],
+                "state": "{{ states('number.test_state') }}",
+                "availability": "{{ is_state('binary_sensor.test_availability', 'on') }}",
+            },
+        )
+    ],
+)
+@pytest.mark.parametrize(
+    "style", [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER]
+)
+@pytest.mark.usefixtures("setup_number")
+async def test_availability(hass: HomeAssistant) -> None:
+    """Test configuration with optimistic state."""
+
+    hass.states.async_set(TEST_AVAILABILITY_ENTITY_ID, "on")
+    hass.states.async_set(TEST_STATE_ENTITY_ID, "4.0")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_NUMBER)
+    assert float(state.state) == 4
+
+    hass.states.async_set(TEST_AVAILABILITY_ENTITY_ID, "off")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_NUMBER)
+    assert state.state == STATE_UNAVAILABLE
+
+    hass.states.async_set(TEST_STATE_ENTITY_ID, "2.0")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_NUMBER)
+    assert state.state == STATE_UNAVAILABLE
+
+    hass.states.async_set(TEST_AVAILABILITY_ENTITY_ID, "on")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_NUMBER)
+    assert float(state.state) == 2
+
+
 @pytest.mark.parametrize(
     ("count", "number_config"),
     [