diff --git a/homeassistant/components/srp_energy/__init__.py b/homeassistant/components/srp_energy/__init__.py index ea80a29d990..98d1cdd421a 100644 --- a/homeassistant/components/srp_energy/__init__.py +++ b/homeassistant/components/srp_energy/__init__.py @@ -5,7 +5,8 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ID, CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN, LOGGER +from .const import CONF_IS_TOU, DOMAIN, LOGGER +from .coordinator import SRPEnergyDataUpdateCoordinator PLATFORMS = [Platform.SENSOR] @@ -24,8 +25,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: api_password, ) + coordinator = SRPEnergyDataUpdateCoordinator( + hass, api_instance, entry.data[CONF_IS_TOU] + ) + + await coordinator.async_config_entry_first_refresh() + hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = api_instance + hass.data[DOMAIN][entry.entry_id] = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) diff --git a/homeassistant/components/srp_energy/coordinator.py b/homeassistant/components/srp_energy/coordinator.py new file mode 100644 index 00000000000..9637f176886 --- /dev/null +++ b/homeassistant/components/srp_energy/coordinator.py @@ -0,0 +1,77 @@ +"""DataUpdateCoordinator for the srp_energy integration.""" +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from requests.exceptions import ConnectionError as ConnectError, HTTPError, Timeout +from srpenergy.client import SrpEnergyClient + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed +from homeassistant.util import dt as dt_util + +from .const import DOMAIN, LOGGER, MIN_TIME_BETWEEN_UPDATES, PHOENIX_TIME_ZONE + + +class SRPEnergyDataUpdateCoordinator(DataUpdateCoordinator[float]): + """A srp_energy Data Update Coordinator.""" + + config_entry: ConfigEntry + + def __init__( + self, hass: HomeAssistant, client: SrpEnergyClient, is_time_of_use: bool + ) -> None: + """Initialize the srp_energy data coordinator.""" + self._client = client + self._is_time_of_use = is_time_of_use + super().__init__( + hass, + LOGGER, + name=DOMAIN, + update_interval=MIN_TIME_BETWEEN_UPDATES, + ) + + async def _async_update_data(self) -> float: + """Fetch data from API endpoint. + + This is the place to pre-process the data to lookup tables + so entities can quickly look up their data. + """ + LOGGER.debug("async_update_data enter") + try: + # Fetch srp_energy data + phx_time_zone = dt_util.get_time_zone(PHOENIX_TIME_ZONE) + end_date = dt_util.now(phx_time_zone) + start_date = end_date - timedelta(days=1) + + async with asyncio.timeout(10): + hourly_usage = await self.hass.async_add_executor_job( + self._client.usage, + start_date, + end_date, + self._is_time_of_use, + ) + + LOGGER.debug( + "async_update_data: Received %s record(s) from %s to %s", + len(hourly_usage) if hourly_usage else "None", + start_date, + end_date, + ) + + previous_daily_usage = 0.0 + for _, _, _, kwh, _ in hourly_usage: + previous_daily_usage += float(kwh) + + LOGGER.debug( + "async_update_data: previous_daily_usage %s", + previous_daily_usage, + ) + + return previous_daily_usage + except TimeoutError as timeout_err: + raise UpdateFailed("Timeout communicating with API") from timeout_err + except (ConnectError, HTTPError, Timeout, ValueError, TypeError) as err: + raise UpdateFailed(f"Error communicating with API: {err}") from err diff --git a/homeassistant/components/srp_energy/sensor.py b/homeassistant/components/srp_energy/sensor.py index d477b65b21d..601baaee8ca 100644 --- a/homeassistant/components/srp_energy/sensor.py +++ b/homeassistant/components/srp_energy/sensor.py @@ -1,11 +1,6 @@ """Support for SRP Energy Sensor.""" from __future__ import annotations -import asyncio -from datetime import timedelta - -from requests.exceptions import ConnectionError as ConnectError, HTTPError, Timeout - from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, @@ -15,88 +10,22 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfEnergy from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.typing import StateType -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed -from homeassistant.util import dt as dt_util +from homeassistant.helpers.update_coordinator import CoordinatorEntity -from .const import ( - CONF_IS_TOU, - DEFAULT_NAME, - DOMAIN, - LOGGER, - MIN_TIME_BETWEEN_UPDATES, - PHOENIX_TIME_ZONE, - SENSOR_NAME, - SENSOR_TYPE, -) +from . import SRPEnergyDataUpdateCoordinator +from .const import DEFAULT_NAME, DOMAIN, SENSOR_NAME async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up the SRP Energy Usage sensor.""" - # API object stored here by __init__.py - api = hass.data[DOMAIN][entry.entry_id] - is_time_of_use = entry.data[CONF_IS_TOU] - - async def async_update_data(): - """Fetch data from API endpoint. - - This is the place to pre-process the data to lookup tables - so entities can quickly look up their data. - """ - LOGGER.debug("async_update_data enter") - try: - # Fetch srp_energy data - phx_time_zone = dt_util.get_time_zone(PHOENIX_TIME_ZONE) - end_date = dt_util.now(phx_time_zone) - start_date = end_date - timedelta(days=1) - - async with asyncio.timeout(10): - hourly_usage = await hass.async_add_executor_job( - api.usage, - start_date, - end_date, - is_time_of_use, - ) - - LOGGER.debug( - "async_update_data: Received %s record(s) from %s to %s", - len(hourly_usage) if hourly_usage else "None", - start_date, - end_date, - ) - - previous_daily_usage = 0.0 - for _, _, _, kwh, _ in hourly_usage: - previous_daily_usage += float(kwh) - - LOGGER.debug( - "async_update_data: previous_daily_usage %s", - previous_daily_usage, - ) - - return previous_daily_usage - except TimeoutError as timeout_err: - raise UpdateFailed("Timeout communicating with API") from timeout_err - except (ConnectError, HTTPError, Timeout, ValueError, TypeError) as err: - raise UpdateFailed(f"Error communicating with API: {err}") from err - - coordinator = DataUpdateCoordinator( - hass, - LOGGER, - name="sensor", - update_method=async_update_data, - update_interval=MIN_TIME_BETWEEN_UPDATES, - ) - - # Fetch initial data so we have data when entities subscribe - await coordinator.async_refresh() + coordinator: SRPEnergyDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] async_add_entities([SrpEntity(coordinator)]) -class SrpEntity(SensorEntity): +class SrpEntity(CoordinatorEntity[SRPEnergyDataUpdateCoordinator], SensorEntity): """Implementation of a Srp Energy Usage sensor.""" _attr_attribution = "Powered by SRP Energy" @@ -104,13 +33,11 @@ class SrpEntity(SensorEntity): _attr_native_unit_of_measurement = UnitOfEnergy.KILO_WATT_HOUR _attr_device_class = SensorDeviceClass.ENERGY _attr_state_class = SensorStateClass.TOTAL_INCREASING - _attr_should_poll = False - def __init__(self, coordinator) -> None: + def __init__(self, coordinator: SRPEnergyDataUpdateCoordinator) -> None: """Initialize the SrpEntity class.""" + super().__init__(coordinator) self._name = SENSOR_NAME - self.type = SENSOR_TYPE - self.coordinator = coordinator @property def name(self) -> str: @@ -118,24 +45,6 @@ class SrpEntity(SensorEntity): return f"{DEFAULT_NAME} {self._name}" @property - def native_value(self) -> StateType: + def native_value(self) -> float: """Return the state of the device.""" return self.coordinator.data - - @property - def available(self) -> bool: - """Return if entity is available.""" - return self.coordinator.last_update_success - - async def async_added_to_hass(self) -> None: - """When entity is added to hass.""" - self.async_on_remove( - self.coordinator.async_add_listener(self.async_write_ha_state) - ) - - async def async_update(self) -> None: - """Update the entity. - - Only used by the generic entity update service. - """ - await self.coordinator.async_request_refresh() diff --git a/tests/components/srp_energy/test_sensor.py b/tests/components/srp_energy/test_sensor.py index 3310e9ce9cd..0e3075c6ac8 100644 --- a/tests/components/srp_energy/test_sensor.py +++ b/tests/components/srp_energy/test_sensor.py @@ -1,4 +1,9 @@ """Tests for the srp_energy sensor platform.""" +from unittest.mock import patch + +import pytest +from requests.models import HTTPError + from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass from homeassistant.config_entries import ConfigEntryState from homeassistant.const import ( @@ -10,6 +15,8 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant +from tests.common import MockConfigEntry + async def test_loading_sensors(hass: HomeAssistant, init_integration) -> None: """Test the srp energy sensors.""" @@ -37,3 +44,26 @@ async def test_srp_entity(hass: HomeAssistant, init_integration) -> None: assert usage_state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.ENERGY assert usage_state.attributes.get(ATTR_ICON) == "mdi:flash" + + +@pytest.mark.parametrize("error", [TimeoutError, HTTPError]) +async def test_srp_entity_update_failed( + hass: HomeAssistant, + error: Exception, + mock_config_entry: MockConfigEntry, +) -> None: + """Test the SrpEntity.""" + + with patch( + "homeassistant.components.srp_energy.SrpEnergyClient", autospec=True + ) as srp_energy_mock: + client = srp_energy_mock.return_value + client.validate.return_value = True + client.usage.side_effect = error + mock_config_entry.add_to_hass(hass) + + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() + + usage_state = hass.states.get("sensor.home_energy_usage") + assert usage_state is None