Use DataUpdateCoordinator in scrape (#80593)
* Add DataUpdateCoordinator to scrape * Fix tests
This commit is contained in:
parent
ebfb10c177
commit
64d6d04ade
3 changed files with 70 additions and 30 deletions
36
homeassistant/components/scrape/coordinator.py
Normal file
36
homeassistant/components/scrape/coordinator.py
Normal 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")
|
|
@ -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()
|
||||||
|
|
|
@ -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:
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue