Refactor Rest Sensor with ManualTriggerEntity (#97396)

* ManualTriggerEntity for rest sensor

* add availability test

* review comments

* last fixes

homeassistant/components/rest/schema.py

@@ -76,6 +76,7 @@ SENSOR_SCHEMA = {
     vol.Optional(CONF_JSON_ATTRS_PATH): cv.string,
     vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
     vol.Optional(CONF_FORCE_UPDATE, default=DEFAULT_FORCE_UPDATE): cv.boolean,
+    vol.Optional(CONF_AVAILABILITY): cv.template,
 }
 
 BINARY_SENSOR_SCHEMA = {

homeassistant/components/rest/sensor.py

@@ -3,28 +3,40 @@ from __future__ import annotations
 
 import logging
 import ssl
+from typing import Any
 
 from jsonpath import jsonpath
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
     DOMAIN as SENSOR_DOMAIN,
     PLATFORM_SCHEMA,
     SensorDeviceClass,
+    SensorEntity,
 )
 from homeassistant.components.sensor.helpers import async_parse_date_datetime
 from homeassistant.const import (
+    CONF_DEVICE_CLASS,
     CONF_FORCE_UPDATE,
+    CONF_ICON,
+    CONF_NAME,
     CONF_RESOURCE,
     CONF_RESOURCE_TEMPLATE,
     CONF_UNIQUE_ID,
+    CONF_UNIT_OF_MEASUREMENT,
     CONF_VALUE_TEMPLATE,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import PlatformNotReady
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.template_entity import TemplateSensor
+from homeassistant.helpers.template import Template
+from homeassistant.helpers.template_entity import (
+    CONF_AVAILABILITY,
+    CONF_PICTURE,
+    ManualTriggerSensorEntity,
+)
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from homeassistant.util.json import json_loads
@@ -43,6 +55,16 @@ PLATFORM_SCHEMA = vol.All(
     cv.has_at_least_one_key(CONF_RESOURCE, CONF_RESOURCE_TEMPLATE), PLATFORM_SCHEMA
 )
 
+TRIGGER_ENTITY_OPTIONS = (
+    CONF_AVAILABILITY,
+    CONF_DEVICE_CLASS,
+    CONF_ICON,
+    CONF_PICTURE,
+    CONF_UNIQUE_ID,
+    CONF_STATE_CLASS,
+    CONF_UNIT_OF_MEASUREMENT,
+)
+
 
 async def async_setup_platform(
     hass: HomeAssistant,
@@ -75,7 +97,14 @@ async def async_setup_platform(
             raise PlatformNotReady from rest.last_exception
         raise PlatformNotReady
 
-    unique_id: str | None = conf.get(CONF_UNIQUE_ID)
+    name = conf.get(CONF_NAME) or Template(DEFAULT_SENSOR_NAME, hass)
+
+    trigger_entity_config = {CONF_NAME: name}
+
+    for key in TRIGGER_ENTITY_OPTIONS:
+        if key not in conf:
+            continue
+        trigger_entity_config[key] = conf[key]
 
     async_add_entities(
         [
@@ -84,13 +113,13 @@ async def async_setup_platform(
                 coordinator,
                 rest,
                 conf,
-                unique_id,
+                trigger_entity_config,
             )
         ],
     )
 
 
-class RestSensor(RestEntity, TemplateSensor):
+class RestSensor(ManualTriggerSensorEntity, RestEntity, SensorEntity):
     """Implementation of a REST sensor."""
 
     def __init__(
@@ -99,9 +128,10 @@ class RestSensor(RestEntity, TemplateSensor):
         coordinator: DataUpdateCoordinator[None] | None,
         rest: RestData,
         config: ConfigType,
-        unique_id: str | None,
+        trigger_entity_config: ConfigType,
     ) -> None:
         """Initialize the REST sensor."""
+        ManualTriggerSensorEntity.__init__(self, hass, trigger_entity_config)
         RestEntity.__init__(
             self,
             coordinator,
@@ -109,25 +139,30 @@ class RestSensor(RestEntity, TemplateSensor):
             config.get(CONF_RESOURCE_TEMPLATE),
             config[CONF_FORCE_UPDATE],
         )
-        TemplateSensor.__init__(
-            self,
-            hass,
-            config=config,
-            fallback_name=DEFAULT_SENSOR_NAME,
-            unique_id=unique_id,
-        )
         self._value_template = config.get(CONF_VALUE_TEMPLATE)
         if (value_template := self._value_template) is not None:
             value_template.hass = hass
         self._json_attrs = config.get(CONF_JSON_ATTRS)
         self._json_attrs_path = config.get(CONF_JSON_ATTRS_PATH)
+        self._attr_extra_state_attributes = {}
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        available1 = RestEntity.available.fget(self)  # type: ignore[attr-defined]
+        available2 = ManualTriggerSensorEntity.available.fget(self)  # type: ignore[attr-defined]
+        return bool(available1 and available2)
+
+    @property
+    def extra_state_attributes(self) -> dict[str, Any]:
+        """Return extra attributes."""
+        return dict(self._attr_extra_state_attributes)
 
     def _update_from_rest_data(self) -> None:
         """Update state from the rest data."""
         value = self.rest.data_without_xml()
 
         if self._json_attrs:
-            self._attr_extra_state_attributes = {}
             if value:
                 try:
                     json_dict = json_loads(value)
@@ -155,6 +190,8 @@ class RestSensor(RestEntity, TemplateSensor):
             else:
                 _LOGGER.warning("Empty reply found when expecting JSON data")
 
+        raw_value = value
+
         if value is not None and self._value_template is not None:
             value = self._value_template.async_render_with_possible_json_value(
                 value, None
@@ -165,8 +202,13 @@ class RestSensor(RestEntity, TemplateSensor):
             SensorDeviceClass.TIMESTAMP,
         ):
             self._attr_native_value = value
+            self._process_manual_data(raw_value)
+            self.async_write_ha_state()
             return
 
         self._attr_native_value = async_parse_date_datetime(
             value, self.entity_id, self.device_class
         )
+
+        self._process_manual_data(raw_value)
+        self.async_write_ha_state()

homeassistant/helpers/template_entity.py

@@ -653,3 +653,17 @@ class ManualTriggerEntity(TriggerBaseEntity):
         variables = {"this": this, **(run_variables or {})}
 
         self._render_templates(variables)
+
+
+class ManualTriggerSensorEntity(ManualTriggerEntity):
+    """Template entity based on manual trigger data for sensor."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config: dict,
+    ) -> None:
+        """Initialize the sensor entity."""
+        ManualTriggerEntity.__init__(self, hass, config)
+        self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
+        self._attr_state_class = config.get(CONF_STATE_CLASS)

tests/components/rest/test_sensor.py

@@ -23,6 +23,7 @@ from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
     CONTENT_TYPE_JSON,
     SERVICE_RELOAD,
+    STATE_UNAVAILABLE,
     STATE_UNKNOWN,
     UnitOfInformation,
     UnitOfTemperature,
@@ -1018,3 +1019,27 @@ async def test_entity_config(hass: HomeAssistant) -> None:
         "state_class": "measurement",
         "unit_of_measurement": "°C",
     }
+
+
+@respx.mock
+async def test_availability_in_config(hass: HomeAssistant) -> None:
+    """Test entity configuration."""
+
+    config = {
+        SENSOR_DOMAIN: {
+            # REST configuration
+            "platform": DOMAIN,
+            "method": "GET",
+            "resource": "http://localhost",
+            # Entity configuration
+            "availability": "{{value==1}}",
+            "name": "{{'REST' + ' ' + 'Sensor'}}",
+        },
+    }
+
+    respx.get("http://localhost").respond(status_code=HTTPStatus.OK, text="123")
+    assert await async_setup_component(hass, SENSOR_DOMAIN, config)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.rest_sensor")
+    assert state.state == STATE_UNAVAILABLE