Scrape refactor to ManualTriggerEntity (#96329)

pull/95672/head
G Johansson 2023-07-25 00:54:19 +02:00 committed by GitHub
parent 5ec633a839
commit c312dcbc4b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 163 additions and 30 deletions

View File

@ -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,

View File

@ -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:

View File

@ -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"