From c312dcbc4b8b9a69ea52398a72db995fb0526587 Mon Sep 17 00:00:00 2001 From: G Johansson Date: Tue, 25 Jul 2023 00:54:19 +0200 Subject: [PATCH] Scrape refactor to ManualTriggerEntity (#96329) --- homeassistant/components/scrape/__init__.py | 6 +- homeassistant/components/scrape/sensor.py | 94 +++++++++++++++------ tests/components/scrape/test_sensor.py | 93 +++++++++++++++++++- 3 files changed, 163 insertions(+), 30 deletions(-) diff --git a/homeassistant/components/scrape/__init__.py b/homeassistant/components/scrape/__init__.py index 8953d9facd0..bf2ccb16b03 100644 --- a/homeassistant/components/scrape/__init__.py +++ b/homeassistant/components/scrape/__init__.py @@ -20,7 +20,10 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers import discovery import homeassistant.helpers.config_validation as cv -from homeassistant.helpers.template_entity import TEMPLATE_SENSOR_BASE_SCHEMA +from homeassistant.helpers.template_entity import ( + CONF_AVAILABILITY, + TEMPLATE_SENSOR_BASE_SCHEMA, +) from homeassistant.helpers.typing import ConfigType from .const import CONF_INDEX, CONF_SELECT, DEFAULT_SCAN_INTERVAL, DOMAIN, PLATFORMS @@ -29,6 +32,7 @@ from .coordinator import ScrapeCoordinator SENSOR_SCHEMA = vol.Schema( { **TEMPLATE_SENSOR_BASE_SCHEMA.schema, + vol.Optional(CONF_AVAILABILITY): cv.template, vol.Optional(CONF_ATTRIBUTE): cv.string, vol.Optional(CONF_INDEX, default=0): cv.positive_int, vol.Required(CONF_SELECT): cv.string, diff --git a/homeassistant/components/scrape/sensor.py b/homeassistant/components/scrape/sensor.py index 5ddd6c48e43..a68083856f7 100644 --- a/homeassistant/components/scrape/sensor.py +++ b/homeassistant/components/scrape/sensor.py @@ -6,13 +6,20 @@ from typing import Any, cast import voluptuous as vol -from homeassistant.components.sensor import SensorDeviceClass +from homeassistant.components.sensor import ( + CONF_STATE_CLASS, + SensorDeviceClass, + SensorEntity, +) from homeassistant.components.sensor.helpers import async_parse_date_datetime from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_ATTRIBUTE, + CONF_DEVICE_CLASS, + CONF_ICON, CONF_NAME, CONF_UNIQUE_ID, + CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE, ) from homeassistant.core import HomeAssistant, callback @@ -20,8 +27,10 @@ from homeassistant.exceptions import PlatformNotReady from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.template import Template from homeassistant.helpers.template_entity import ( + CONF_AVAILABILITY, + CONF_PICTURE, TEMPLATE_SENSOR_BASE_SCHEMA, - TemplateSensor, + ManualTriggerEntity, ) from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType from homeassistant.helpers.update_coordinator import CoordinatorEntity @@ -53,17 +62,30 @@ async def async_setup_platform( if value_template is not None: value_template.hass = hass + trigger_entity_config = { + CONF_NAME: sensor_config[CONF_NAME], + CONF_DEVICE_CLASS: sensor_config.get(CONF_DEVICE_CLASS), + CONF_UNIQUE_ID: sensor_config.get(CONF_UNIQUE_ID), + } + if available := sensor_config.get(CONF_AVAILABILITY): + trigger_entity_config[CONF_AVAILABILITY] = available + if icon := sensor_config.get(CONF_ICON): + trigger_entity_config[CONF_ICON] = icon + if picture := sensor_config.get(CONF_PICTURE): + trigger_entity_config[CONF_PICTURE] = picture + entities.append( ScrapeSensor( hass, coordinator, - sensor_config, - sensor_config[CONF_NAME], - sensor_config.get(CONF_UNIQUE_ID), + trigger_entity_config, + sensor_config.get(CONF_UNIT_OF_MEASUREMENT), + sensor_config.get(CONF_STATE_CLASS), sensor_config[CONF_SELECT], sensor_config.get(CONF_ATTRIBUTE), sensor_config[CONF_INDEX], value_template, + True, ) ) @@ -84,60 +106,65 @@ async def async_setup_entry( )(sensor) name: str = sensor_config[CONF_NAME] - select: str = sensor_config[CONF_SELECT] - attr: str | None = sensor_config.get(CONF_ATTRIBUTE) - index: int = int(sensor_config[CONF_INDEX]) value_string: str | None = sensor_config.get(CONF_VALUE_TEMPLATE) - unique_id: str = sensor_config[CONF_UNIQUE_ID] value_template: Template | None = ( Template(value_string, hass) if value_string is not None else None ) + + trigger_entity_config = { + CONF_NAME: name, + CONF_DEVICE_CLASS: sensor_config.get(CONF_DEVICE_CLASS), + CONF_UNIQUE_ID: sensor_config[CONF_UNIQUE_ID], + } + entities.append( ScrapeSensor( hass, coordinator, - sensor_config, - name, - unique_id, - select, - attr, - index, + trigger_entity_config, + sensor_config.get(CONF_UNIT_OF_MEASUREMENT), + sensor_config.get(CONF_STATE_CLASS), + sensor_config[CONF_SELECT], + sensor_config.get(CONF_ATTRIBUTE), + sensor_config[CONF_INDEX], value_template, + False, ) ) async_add_entities(entities) -class ScrapeSensor(CoordinatorEntity[ScrapeCoordinator], TemplateSensor): +class ScrapeSensor( + CoordinatorEntity[ScrapeCoordinator], ManualTriggerEntity, SensorEntity +): """Representation of a web scrape sensor.""" def __init__( self, hass: HomeAssistant, coordinator: ScrapeCoordinator, - config: ConfigType, - name: str, - unique_id: str | None, + trigger_entity_config: ConfigType, + unit_of_measurement: str | None, + state_class: str | None, select: str, attr: str | None, index: int, value_template: Template | None, + yaml: bool, ) -> None: """Initialize a web scrape sensor.""" CoordinatorEntity.__init__(self, coordinator) - TemplateSensor.__init__( - self, - hass, - config=config, - fallback_name=name, - unique_id=unique_id, - ) + ManualTriggerEntity.__init__(self, hass, trigger_entity_config) + self._attr_name = trigger_entity_config[CONF_NAME].template + self._attr_native_unit_of_measurement = unit_of_measurement + self._attr_state_class = state_class self._select = select self._attr = attr self._index = index self._value_template = value_template + self._attr_native_value = None def _extract_value(self) -> Any: """Parse the html extraction in the executor.""" @@ -164,12 +191,15 @@ class ScrapeSensor(CoordinatorEntity[ScrapeCoordinator], TemplateSensor): async def async_added_to_hass(self) -> None: """Ensure the data from the initial update is reflected in the state.""" - await super().async_added_to_hass() + await ManualTriggerEntity.async_added_to_hass(self) + # https://github.com/python/mypy/issues/15097 + await CoordinatorEntity.async_added_to_hass(self) # type: ignore[arg-type] self._async_update_from_rest_data() def _async_update_from_rest_data(self) -> None: """Update state from the rest data.""" value = self._extract_value() + raw_value = value if (template := self._value_template) is not None: value = template.async_render_with_possible_json_value(value, None) @@ -179,11 +209,21 @@ class ScrapeSensor(CoordinatorEntity[ScrapeCoordinator], TemplateSensor): SensorDeviceClass.TIMESTAMP, }: self._attr_native_value = value + self._process_manual_data(raw_value) 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() + + @property + def available(self) -> bool: + """Return if entity is available.""" + available1 = CoordinatorEntity.available.fget(self) # type: ignore[attr-defined] + available2 = ManualTriggerEntity.available.fget(self) # type: ignore[attr-defined] + return bool(available1 and available2) @callback def _handle_coordinator_update(self) -> None: diff --git a/tests/components/scrape/test_sensor.py b/tests/components/scrape/test_sensor.py index 44c264520d6..60cde48e5bf 100644 --- a/tests/components/scrape/test_sensor.py +++ b/tests/components/scrape/test_sensor.py @@ -1,12 +1,16 @@ """The tests for the Scrape sensor platform.""" from __future__ import annotations -from datetime import datetime +from datetime import datetime, timedelta from unittest.mock import patch import pytest -from homeassistant.components.scrape.const import DEFAULT_SCAN_INTERVAL +from homeassistant.components.scrape.const import ( + CONF_INDEX, + CONF_SELECT, + DEFAULT_SCAN_INTERVAL, +) from homeassistant.components.sensor import ( CONF_STATE_CLASS, SensorDeviceClass, @@ -14,6 +18,9 @@ from homeassistant.components.sensor import ( ) from homeassistant.const import ( CONF_DEVICE_CLASS, + CONF_ICON, + CONF_NAME, + CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, STATE_UNAVAILABLE, STATE_UNKNOWN, @@ -21,7 +28,9 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.template_entity import CONF_AVAILABILITY, CONF_PICTURE from homeassistant.setup import async_setup_component +from homeassistant.util import dt as dt_util from . import MockRestData, return_integration_config @@ -469,3 +478,83 @@ async def test_setup_config_entry( entity = entity_reg.async_get("sensor.current_version") assert entity.unique_id == "3699ef88-69e6-11ed-a1eb-0242ac120002" + + +async def test_templates_with_yaml(hass: HomeAssistant) -> None: + """Test the Scrape sensor from yaml config with templates.""" + + hass.states.async_set("sensor.input1", "on") + hass.states.async_set("sensor.input2", "on") + await hass.async_block_till_done() + + config = { + DOMAIN: [ + return_integration_config( + sensors=[ + { + CONF_NAME: "Get values with template", + CONF_SELECT: ".current-version h1", + CONF_INDEX: 0, + CONF_UNIQUE_ID: "3699ef88-69e6-11ed-a1eb-0242ac120002", + CONF_ICON: '{% if states("sensor.input1")=="on" %} mdi:on {% else %} mdi:off {% endif %}', + CONF_PICTURE: '{% if states("sensor.input1")=="on" %} /local/picture1.jpg {% else %} /local/picture2.jpg {% endif %}', + CONF_AVAILABILITY: '{{ states("sensor.input2")=="on" }}', + } + ] + ) + ] + } + + mocker = MockRestData("test_scrape_sensor") + with patch( + "homeassistant.components.rest.RestData", + return_value=mocker, + ): + assert await async_setup_component(hass, DOMAIN, config) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == "Current Version: 2021.12.10" + assert state.attributes[CONF_ICON] == "mdi:on" + assert state.attributes["entity_picture"] == "/local/picture1.jpg" + + hass.states.async_set("sensor.input1", "off") + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(minutes=10), + ) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == "Current Version: 2021.12.10" + assert state.attributes[CONF_ICON] == "mdi:off" + assert state.attributes["entity_picture"] == "/local/picture2.jpg" + + hass.states.async_set("sensor.input2", "off") + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(minutes=20), + ) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == STATE_UNAVAILABLE + + hass.states.async_set("sensor.input1", "on") + hass.states.async_set("sensor.input2", "on") + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(minutes=30), + ) + await hass.async_block_till_done() + + state = hass.states.get("sensor.get_values_with_template") + assert state.state == "Current Version: 2021.12.10" + assert state.attributes[CONF_ICON] == "mdi:on" + assert state.attributes["entity_picture"] == "/local/picture1.jpg"