1
mirror of https://github.com/home-assistant/core synced 2024-07-30 21:18:57 +02:00

Use DataUpdateCoordinator in scrape (#80593)

* Add DataUpdateCoordinator to scrape

* Fix tests
This commit is contained in:
epenet 2022-10-24 14:55:57 +02:00 committed by GitHub
parent ebfb10c177
commit 64d6d04ade
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 70 additions and 30 deletions

View File

@ -0,0 +1,36 @@
"""Coordinator for the scrape component."""
from __future__ import annotations
from datetime import timedelta
import logging
from bs4 import BeautifulSoup
from homeassistant.components.rest.data import RestData
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
_LOGGER = logging.getLogger(__name__)
class ScrapeCoordinator(DataUpdateCoordinator[BeautifulSoup]):
"""Scrape Coordinator."""
def __init__(
self, hass: HomeAssistant, rest: RestData, update_interval: timedelta
) -> None:
"""Initialize Scrape coordinator."""
super().__init__(
hass,
_LOGGER,
name="Scrape Coordinator",
update_interval=update_interval,
)
self._rest = rest
async def _async_update_data(self) -> BeautifulSoup:
"""Fetch data from Rest."""
await self._rest.async_update()
if (data := self._rest.data) is None:
raise UpdateFailed("REST data is not available")
return await self.hass.async_add_executor_job(BeautifulSoup, data, "lxml")

View File

@ -5,7 +5,6 @@ from datetime import timedelta
import logging import logging
from typing import Any from typing import Any
from bs4 import BeautifulSoup
import httpx import httpx
import voluptuous as vol import voluptuous as vol
@ -31,12 +30,15 @@ from homeassistant.const import (
HTTP_BASIC_AUTHENTICATION, HTTP_BASIC_AUTHENTICATION,
HTTP_DIGEST_AUTHENTICATION, HTTP_DIGEST_AUTHENTICATION,
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import PlatformNotReady from homeassistant.exceptions import PlatformNotReady
import homeassistant.helpers.config_validation as cv import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.template import Template from homeassistant.helpers.template import Template
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .coordinator import ScrapeCoordinator
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -105,15 +107,16 @@ async def async_setup_platform(
auth = (username, password) auth = (username, password)
rest = RestData(hass, method, resource, auth, headers, None, payload, verify_ssl) rest = RestData(hass, method, resource, auth, headers, None, payload, verify_ssl)
await rest.async_update()
if rest.data is None: coordinator = ScrapeCoordinator(hass, rest, SCAN_INTERVAL)
await coordinator.async_refresh()
if coordinator.data is None:
raise PlatformNotReady raise PlatformNotReady
async_add_entities( async_add_entities(
[ [
ScrapeSensor( ScrapeSensor(
rest, coordinator,
name, name,
select, select,
attr, attr,
@ -124,16 +127,15 @@ async def async_setup_platform(
state_class, state_class,
) )
], ],
True,
) )
class ScrapeSensor(SensorEntity): class ScrapeSensor(CoordinatorEntity[ScrapeCoordinator], SensorEntity):
"""Representation of a web scrape sensor.""" """Representation of a web scrape sensor."""
def __init__( def __init__(
self, self,
rest: RestData, coordinator: ScrapeCoordinator,
name: str, name: str,
select: str | None, select: str | None,
attr: str | None, attr: str | None,
@ -144,7 +146,7 @@ class ScrapeSensor(SensorEntity):
state_class: str | None, state_class: str | None,
) -> None: ) -> None:
"""Initialize a web scrape sensor.""" """Initialize a web scrape sensor."""
self.rest = rest super().__init__(coordinator)
self._attr_native_value = None self._attr_native_value = None
self._select = select self._select = select
self._attr = attr self._attr = attr
@ -157,9 +159,8 @@ class ScrapeSensor(SensorEntity):
def _extract_value(self) -> Any: def _extract_value(self) -> Any:
"""Parse the html extraction in the executor.""" """Parse the html extraction in the executor."""
raw_data = BeautifulSoup(self.rest.data, "lxml") raw_data = self.coordinator.data
_LOGGER.debug(raw_data) _LOGGER.debug("Raw beautiful soup: %s", raw_data)
try: try:
if self._attr is not None: if self._attr is not None:
value = raw_data.select(self._select)[self._index][self._attr] value = raw_data.select(self._select)[self._index][self._attr]
@ -177,25 +178,17 @@ class ScrapeSensor(SensorEntity):
"Attribute '%s' not found in %s", self._attr, self.entity_id "Attribute '%s' not found in %s", self._attr, self.entity_id
) )
value = None value = None
_LOGGER.debug(value) _LOGGER.debug("Parsed value: %s", value)
return value return value
async def async_update(self) -> None:
"""Get the latest data from the source and updates the state."""
await self.rest.async_update()
await self._async_update_from_rest_data()
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:
"""Ensure the data from the initial update is reflected in the state.""" """Ensure the data from the initial update is reflected in the state."""
await self._async_update_from_rest_data() await super().async_added_to_hass()
self._async_update_from_rest_data()
async def _async_update_from_rest_data(self) -> None: def _async_update_from_rest_data(self) -> None:
"""Update state from the rest data.""" """Update state from the rest data."""
if self.rest.data is None: value = self._extract_value()
_LOGGER.error("Unable to retrieve data for %s", self.name)
return
value = await self.hass.async_add_executor_job(self._extract_value)
if self._value_template is not None: if self._value_template is not None:
self._attr_native_value = ( self._attr_native_value = (
@ -203,3 +196,9 @@ class ScrapeSensor(SensorEntity):
) )
else: else:
self._attr_native_value = value self._attr_native_value = value
@callback
def _handle_coordinator_update(self) -> None:
"""Handle updated data from the coordinator."""
self._async_update_from_rest_data()
super()._handle_coordinator_update()

View File

@ -1,8 +1,10 @@
"""The tests for the Scrape sensor platform.""" """The tests for the Scrape sensor platform."""
from __future__ import annotations from __future__ import annotations
from datetime import datetime
from unittest.mock import patch from unittest.mock import patch
from homeassistant.components.scrape.sensor import SCAN_INTERVAL
from homeassistant.components.sensor import ( from homeassistant.components.sensor import (
CONF_STATE_CLASS, CONF_STATE_CLASS,
SensorDeviceClass, SensorDeviceClass,
@ -11,15 +13,17 @@ from homeassistant.components.sensor import (
from homeassistant.const import ( from homeassistant.const import (
CONF_DEVICE_CLASS, CONF_DEVICE_CLASS,
CONF_UNIT_OF_MEASUREMENT, CONF_UNIT_OF_MEASUREMENT,
STATE_UNAVAILABLE,
STATE_UNKNOWN, STATE_UNKNOWN,
TEMP_CELSIUS, TEMP_CELSIUS,
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_component import async_update_entity
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
from . import MockRestData, return_config from . import MockRestData, return_config
from tests.common import async_fire_time_changed
DOMAIN = "scrape" DOMAIN = "scrape"
@ -155,12 +159,13 @@ async def test_scrape_sensor_no_data_refresh(hass: HomeAssistant) -> None:
assert state assert state
assert state.state == "Current Version: 2021.12.10" assert state.state == "Current Version: 2021.12.10"
mocker.data = None mocker.payload = "test_scrape_sensor_no_data"
await async_update_entity(hass, "sensor.ha_version") async_fire_time_changed(hass, datetime.utcnow() + SCAN_INTERVAL)
await hass.async_block_till_done()
assert mocker.data is None state = hass.states.get("sensor.ha_version")
assert state is not None assert state is not None
assert state.state == "Current Version: 2021.12.10" assert state.state == STATE_UNAVAILABLE
async def test_scrape_sensor_attribute_and_tag(hass: HomeAssistant) -> None: async def test_scrape_sensor_attribute_and_tag(hass: HomeAssistant) -> None: