diff --git a/homeassistant/components/senz/__init__.py b/homeassistant/components/senz/__init__.py index a23ff7bb994..ac3c1949c34 100644 --- a/homeassistant/components/senz/__init__.py +++ b/homeassistant/components/senz/__init__.py @@ -2,16 +2,14 @@ from __future__ import annotations -from datetime import timedelta from http import HTTPStatus import logging from aiohttp import ClientResponseError from httpx import HTTPStatusError, RequestError import jwt -from pysenz import SENZAPI, Thermostat +from pysenz import SENZAPI -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady @@ -21,12 +19,10 @@ from homeassistant.helpers.config_entry_oauth2_flow import ( OAuth2Session, async_get_config_entry_implementation, ) -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .api import SENZConfigEntryAuth from .const import DOMAIN - -UPDATE_INTERVAL = timedelta(seconds=30) +from .coordinator import SENZConfigEntry, SENZDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -34,9 +30,6 @@ CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) PLATFORMS = [Platform.CLIMATE, Platform.SENSOR] -type SENZDataUpdateCoordinator = DataUpdateCoordinator[dict[str, Thermostat]] -type SENZConfigEntry = ConfigEntry[SENZDataUpdateCoordinator] - async def async_setup_entry(hass: HomeAssistant, entry: SENZConfigEntry) -> bool: """Set up SENZ from a config entry.""" @@ -51,14 +44,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: SENZConfigEntry) -> bool auth = SENZConfigEntryAuth(httpx_client.get_async_client(hass), session) senz_api = SENZAPI(auth) - async def update_thermostats() -> dict[str, Thermostat]: - """Fetch SENZ thermostats data.""" - try: - thermostats = await senz_api.get_thermostats() - except RequestError as err: - raise UpdateFailed from err - return {thermostat.serial_number: thermostat for thermostat in thermostats} - try: account = await senz_api.get_account() except HTTPStatusError as err: @@ -92,13 +77,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: SENZConfigEntry) -> bool translation_key="config_entry_auth_failed", ) from err - coordinator: SENZDataUpdateCoordinator = DataUpdateCoordinator( + coordinator = SENZDataUpdateCoordinator( hass, - _LOGGER, - config_entry=entry, + entry, name=account.username, - update_interval=UPDATE_INTERVAL, - update_method=update_thermostats, + senz_api=senz_api, ) await coordinator.async_config_entry_first_refresh() diff --git a/homeassistant/components/senz/climate.py b/homeassistant/components/senz/climate.py index bde683d60d9..9f5bc15e5bf 100644 --- a/homeassistant/components/senz/climate.py +++ b/homeassistant/components/senz/climate.py @@ -20,8 +20,8 @@ from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import SENZConfigEntry, SENZDataUpdateCoordinator from .const import DOMAIN +from .coordinator import SENZConfigEntry, SENZDataUpdateCoordinator async def async_setup_entry( diff --git a/homeassistant/components/senz/coordinator.py b/homeassistant/components/senz/coordinator.py new file mode 100644 index 00000000000..44f218d7b40 --- /dev/null +++ b/homeassistant/components/senz/coordinator.py @@ -0,0 +1,51 @@ +"""Data update coordinator for SENZ.""" + +from __future__ import annotations + +from datetime import timedelta +import logging + +from httpx import RequestError +from pysenz import SENZAPI, Thermostat + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +UPDATE_INTERVAL = timedelta(seconds=30) + +_LOGGER = logging.getLogger(__name__) + +type SENZConfigEntry = ConfigEntry[SENZDataUpdateCoordinator] + + +class SENZDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Thermostat]]): + """Class to manage fetching SENZ data.""" + + config_entry: SENZConfigEntry + + def __init__( + self, + hass: HomeAssistant, + config_entry: SENZConfigEntry, + *, + name: str, + senz_api: SENZAPI, + ) -> None: + """Initialize the coordinator.""" + super().__init__( + hass, + _LOGGER, + config_entry=config_entry, + name=name, + update_interval=UPDATE_INTERVAL, + ) + self._senz_api = senz_api + + async def _async_update_data(self) -> dict[str, Thermostat]: + """Fetch data from SENZ.""" + try: + thermostats = await self._senz_api.get_thermostats() + except RequestError as err: + raise UpdateFailed from err + return {thermostat.serial_number: thermostat for thermostat in thermostats} diff --git a/homeassistant/components/senz/diagnostics.py b/homeassistant/components/senz/diagnostics.py index bed15a7091a..909ee161998 100644 --- a/homeassistant/components/senz/diagnostics.py +++ b/homeassistant/components/senz/diagnostics.py @@ -5,7 +5,7 @@ from typing import Any from homeassistant.components.diagnostics import async_redact_data from homeassistant.core import HomeAssistant -from . import SENZConfigEntry +from .coordinator import SENZConfigEntry TO_REDACT = [ "access_token", diff --git a/homeassistant/components/senz/sensor.py b/homeassistant/components/senz/sensor.py index 74acf101ae0..8f7eb2cc0eb 100644 --- a/homeassistant/components/senz/sensor.py +++ b/homeassistant/components/senz/sensor.py @@ -19,8 +19,8 @@ from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import SENZConfigEntry, SENZDataUpdateCoordinator from .const import DOMAIN +from .coordinator import SENZConfigEntry, SENZDataUpdateCoordinator @dataclass(kw_only=True, frozen=True) diff --git a/tests/components/senz/test_climate.py b/tests/components/senz/test_climate.py index cc104b744db..c6f2f747a44 100644 --- a/tests/components/senz/test_climate.py +++ b/tests/components/senz/test_climate.py @@ -52,7 +52,8 @@ async def test_set_target( with ( patch("homeassistant.components.senz.PLATFORMS", [Platform.CLIMATE]), patch( - "homeassistant.components.senz.Thermostat.manual", return_value=None + "homeassistant.components.senz.coordinator.Thermostat.manual", + return_value=None, ) as mock_manual, ): await setup_integration(hass, mock_config_entry) @@ -75,7 +76,7 @@ async def test_set_target_fail( with ( patch("homeassistant.components.senz.PLATFORMS", [Platform.CLIMATE]), patch( - "homeassistant.components.senz.Thermostat.manual", + "homeassistant.components.senz.coordinator.Thermostat.manual", side_effect=RequestError("API error"), ) as mock_manual, ): @@ -109,10 +110,12 @@ async def test_set_hvac_mode( with ( patch("homeassistant.components.senz.PLATFORMS", [Platform.CLIMATE]), patch( - "homeassistant.components.senz.Thermostat.manual", return_value=None + "homeassistant.components.senz.coordinator.Thermostat.manual", + return_value=None, ) as mock_manual, patch( - "homeassistant.components.senz.Thermostat.auto", return_value=None + "homeassistant.components.senz.coordinator.Thermostat.auto", + return_value=None, ) as mock_auto, ): await setup_integration(hass, mock_config_entry) @@ -136,7 +139,7 @@ async def test_set_hvac_mode_fail( with ( patch("homeassistant.components.senz.PLATFORMS", [Platform.CLIMATE]), patch( - "homeassistant.components.senz.Thermostat.manual", + "homeassistant.components.senz.coordinator.Thermostat.manual", side_effect=RequestError("API error"), ) as mock_manual, ):