From 09483d2cef3e90348cde39564a1263d7f86f5ec0 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Wed, 5 Feb 2025 18:27:13 +0100 Subject: [PATCH] Move fireservicerota coordinator to own module (#137304) --- .../components/fireservicerota/__init__.py | 193 +--------------- .../fireservicerota/binary_sensor.py | 15 +- .../components/fireservicerota/coordinator.py | 213 ++++++++++++++++++ .../fireservicerota/test_config_flow.py | 2 +- 4 files changed, 225 insertions(+), 198 deletions(-) create mode 100644 homeassistant/components/fireservicerota/coordinator.py diff --git a/homeassistant/components/fireservicerota/__init__.py b/homeassistant/components/fireservicerota/__init__.py index aa303a08795..360a0f0b210 100644 --- a/homeassistant/components/fireservicerota/__init__.py +++ b/homeassistant/components/fireservicerota/__init__.py @@ -3,29 +3,16 @@ from __future__ import annotations from datetime import timedelta -import logging - -from pyfireservicerota import ( - ExpiredTokenError, - FireServiceRota, - FireServiceRotaIncidents, - InvalidAuthError, - InvalidTokenError, -) from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_TOKEN, CONF_URL, CONF_USERNAME, Platform +from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from homeassistant.exceptions import ConfigEntryAuthFailed -from homeassistant.helpers.dispatcher import dispatcher_send -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN, WSS_BWRURL +from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN +from .coordinator import FireServiceRotaClient, FireServiceUpdateCoordinator MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=60) -_LOGGER = logging.getLogger(__name__) - PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH] @@ -40,17 +27,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if client.token_refresh_failure: return False - async def async_update_data(): - return await client.async_update() - - coordinator = DataUpdateCoordinator( - hass, - _LOGGER, - config_entry=entry, - name="duty binary sensor", - update_method=async_update_data, - update_interval=MIN_TIME_BETWEEN_UPDATES, - ) + coordinator = FireServiceUpdateCoordinator(hass, client, entry) await coordinator.async_config_entry_first_refresh() @@ -74,165 +51,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if unload_ok: del hass.data[DOMAIN][entry.entry_id] return unload_ok - - -class FireServiceRotaOauth: - """Handle authentication tokens.""" - - def __init__(self, hass, entry, fsr): - """Initialize the oauth object.""" - self._hass = hass - self._entry = entry - - self._url = entry.data[CONF_URL] - self._username = entry.data[CONF_USERNAME] - self._fsr = fsr - - async def async_refresh_tokens(self) -> bool: - """Refresh tokens and update config entry.""" - _LOGGER.debug("Refreshing authentication tokens after expiration") - - try: - token_info = await self._hass.async_add_executor_job( - self._fsr.refresh_tokens - ) - - except (InvalidAuthError, InvalidTokenError) as err: - raise ConfigEntryAuthFailed( - "Error refreshing tokens, triggered reauth workflow" - ) from err - - _LOGGER.debug("Saving new tokens in config entry") - self._hass.config_entries.async_update_entry( - self._entry, - data={ - "auth_implementation": DOMAIN, - CONF_URL: self._url, - CONF_USERNAME: self._username, - CONF_TOKEN: token_info, - }, - ) - - return True - - -class FireServiceRotaWebSocket: - """Define a FireServiceRota websocket manager object.""" - - def __init__(self, hass, entry): - """Initialize the websocket object.""" - self._hass = hass - self._entry = entry - - self._fsr_incidents = FireServiceRotaIncidents(on_incident=self._on_incident) - self.incident_data = None - - def _construct_url(self) -> str: - """Return URL with latest access token.""" - return WSS_BWRURL.format( - self._entry.data[CONF_URL], self._entry.data[CONF_TOKEN]["access_token"] - ) - - def _on_incident(self, data) -> None: - """Received new incident, update data.""" - _LOGGER.debug("Received new incident via websocket: %s", data) - self.incident_data = data - dispatcher_send(self._hass, f"{DOMAIN}_{self._entry.entry_id}_update") - - def start_listener(self) -> None: - """Start the websocket listener.""" - _LOGGER.debug("Starting incidents listener") - self._fsr_incidents.start(self._construct_url()) - - def stop_listener(self) -> None: - """Stop the websocket listener.""" - _LOGGER.debug("Stopping incidents listener") - self._fsr_incidents.stop() - - -class FireServiceRotaClient: - """Getting the latest data from fireservicerota.""" - - def __init__(self, hass, entry): - """Initialize the data object.""" - self._hass = hass - self._entry = entry - - self._url = entry.data[CONF_URL] - self._tokens = entry.data[CONF_TOKEN] - - self.entry_id = entry.entry_id - self.unique_id = entry.unique_id - - self.token_refresh_failure = False - self.incident_id = None - self.on_duty = False - - self.fsr = FireServiceRota(base_url=self._url, token_info=self._tokens) - - self.oauth = FireServiceRotaOauth( - self._hass, - self._entry, - self.fsr, - ) - - self.websocket = FireServiceRotaWebSocket(self._hass, self._entry) - - async def setup(self) -> None: - """Set up the data client.""" - await self._hass.async_add_executor_job(self.websocket.start_listener) - - async def update_call(self, func, *args): - """Perform update call and return data.""" - if self.token_refresh_failure: - return None - - try: - return await self._hass.async_add_executor_job(func, *args) - except (ExpiredTokenError, InvalidTokenError): - await self._hass.async_add_executor_job(self.websocket.stop_listener) - self.token_refresh_failure = True - - if await self.oauth.async_refresh_tokens(): - self.token_refresh_failure = False - await self._hass.async_add_executor_job(self.websocket.start_listener) - - return await self._hass.async_add_executor_job(func, *args) - - async def async_update(self) -> dict | None: - """Get the latest availability data.""" - data = await self.update_call( - self.fsr.get_availability, str(self._hass.config.time_zone) - ) - - if not data: - return None - - self.on_duty = bool(data.get("available")) - - _LOGGER.debug("Updated availability data: %s", data) - return data - - async def async_response_update(self) -> dict | None: - """Get the latest incident response data.""" - - if not self.incident_id: - return None - - _LOGGER.debug("Updating response data for incident id %s", self.incident_id) - - return await self.update_call(self.fsr.get_incident_response, self.incident_id) - - async def async_set_response(self, value) -> None: - """Set incident response status.""" - - if not self.incident_id: - return - - _LOGGER.debug( - "Setting incident response for incident id '%s' to state '%s'", - self.incident_id, - value, - ) - - await self.update_call(self.fsr.set_incident_response, self.incident_id, value) diff --git a/homeassistant/components/fireservicerota/binary_sensor.py b/homeassistant/components/fireservicerota/binary_sensor.py index a22991f2008..b6d3aa67a0a 100644 --- a/homeassistant/components/fireservicerota/binary_sensor.py +++ b/homeassistant/components/fireservicerota/binary_sensor.py @@ -8,13 +8,10 @@ from homeassistant.components.binary_sensor import BinarySensorEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import FireServiceRotaClient from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN as FIRESERVICEROTA_DOMAIN +from .coordinator import FireServiceRotaClient, FireServiceUpdateCoordinator async def async_setup_entry( @@ -26,14 +23,16 @@ async def async_setup_entry( DATA_CLIENT ] - coordinator: DataUpdateCoordinator = hass.data[FIRESERVICEROTA_DOMAIN][ + coordinator: FireServiceUpdateCoordinator = hass.data[FIRESERVICEROTA_DOMAIN][ entry.entry_id ][DATA_COORDINATOR] async_add_entities([ResponseBinarySensor(coordinator, client, entry)]) -class ResponseBinarySensor(CoordinatorEntity, BinarySensorEntity): +class ResponseBinarySensor( + CoordinatorEntity[FireServiceUpdateCoordinator], BinarySensorEntity +): """Representation of an FireServiceRota sensor.""" _attr_has_entity_name = True @@ -41,7 +40,7 @@ class ResponseBinarySensor(CoordinatorEntity, BinarySensorEntity): def __init__( self, - coordinator: DataUpdateCoordinator, + coordinator: FireServiceUpdateCoordinator, client: FireServiceRotaClient, entry: ConfigEntry, ) -> None: diff --git a/homeassistant/components/fireservicerota/coordinator.py b/homeassistant/components/fireservicerota/coordinator.py new file mode 100644 index 00000000000..35f839b3bdb --- /dev/null +++ b/homeassistant/components/fireservicerota/coordinator.py @@ -0,0 +1,213 @@ +"""The FireServiceRota integration.""" + +from __future__ import annotations + +from datetime import timedelta +import logging + +from pyfireservicerota import ( + ExpiredTokenError, + FireServiceRota, + FireServiceRotaIncidents, + InvalidAuthError, + InvalidTokenError, +) + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import CONF_TOKEN, CONF_URL, CONF_USERNAME, Platform +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ConfigEntryAuthFailed +from homeassistant.helpers.dispatcher import dispatcher_send +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator + +from .const import DOMAIN, WSS_BWRURL + +MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=60) + +_LOGGER = logging.getLogger(__name__) + +PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH] + + +class FireServiceUpdateCoordinator(DataUpdateCoordinator[dict | None]): + """Data update coordinator for FireServiceRota.""" + + def __init__( + self, hass: HomeAssistant, client: FireServiceRotaClient, entry: ConfigEntry + ) -> None: + """Initialize the FireServiceRota DataUpdateCoordinator.""" + super().__init__( + hass, + _LOGGER, + name="duty binary sensor", + config_entry=entry, + update_interval=MIN_TIME_BETWEEN_UPDATES, + ) + + self.client = client + + async def _async_update_data(self) -> dict | None: + """Get the latest availability data.""" + return await self.client.async_update() + + +class FireServiceRotaOauth: + """Handle authentication tokens.""" + + def __init__(self, hass, entry, fsr): + """Initialize the oauth object.""" + self._hass = hass + self._entry = entry + + self._url = entry.data[CONF_URL] + self._username = entry.data[CONF_USERNAME] + self._fsr = fsr + + async def async_refresh_tokens(self) -> bool: + """Refresh tokens and update config entry.""" + _LOGGER.debug("Refreshing authentication tokens after expiration") + + try: + token_info = await self._hass.async_add_executor_job( + self._fsr.refresh_tokens + ) + + except (InvalidAuthError, InvalidTokenError) as err: + raise ConfigEntryAuthFailed( + "Error refreshing tokens, triggered reauth workflow" + ) from err + + _LOGGER.debug("Saving new tokens in config entry") + self._hass.config_entries.async_update_entry( + self._entry, + data={ + "auth_implementation": DOMAIN, + CONF_URL: self._url, + CONF_USERNAME: self._username, + CONF_TOKEN: token_info, + }, + ) + + return True + + +class FireServiceRotaWebSocket: + """Define a FireServiceRota websocket manager object.""" + + def __init__(self, hass, entry): + """Initialize the websocket object.""" + self._hass = hass + self._entry = entry + + self._fsr_incidents = FireServiceRotaIncidents(on_incident=self._on_incident) + self.incident_data = None + + def _construct_url(self) -> str: + """Return URL with latest access token.""" + return WSS_BWRURL.format( + self._entry.data[CONF_URL], self._entry.data[CONF_TOKEN]["access_token"] + ) + + def _on_incident(self, data) -> None: + """Received new incident, update data.""" + _LOGGER.debug("Received new incident via websocket: %s", data) + self.incident_data = data + dispatcher_send(self._hass, f"{DOMAIN}_{self._entry.entry_id}_update") + + def start_listener(self) -> None: + """Start the websocket listener.""" + _LOGGER.debug("Starting incidents listener") + self._fsr_incidents.start(self._construct_url()) + + def stop_listener(self) -> None: + """Stop the websocket listener.""" + _LOGGER.debug("Stopping incidents listener") + self._fsr_incidents.stop() + + +class FireServiceRotaClient: + """Getting the latest data from fireservicerota.""" + + def __init__(self, hass, entry): + """Initialize the data object.""" + self._hass = hass + self._entry = entry + + self._url = entry.data[CONF_URL] + self._tokens = entry.data[CONF_TOKEN] + + self.entry_id = entry.entry_id + self.unique_id = entry.unique_id + + self.token_refresh_failure = False + self.incident_id = None + self.on_duty = False + + self.fsr = FireServiceRota(base_url=self._url, token_info=self._tokens) + + self.oauth = FireServiceRotaOauth( + self._hass, + self._entry, + self.fsr, + ) + + self.websocket = FireServiceRotaWebSocket(self._hass, self._entry) + + async def setup(self) -> None: + """Set up the data client.""" + await self._hass.async_add_executor_job(self.websocket.start_listener) + + async def update_call(self, func, *args): + """Perform update call and return data.""" + if self.token_refresh_failure: + return None + + try: + return await self._hass.async_add_executor_job(func, *args) + except (ExpiredTokenError, InvalidTokenError): + await self._hass.async_add_executor_job(self.websocket.stop_listener) + self.token_refresh_failure = True + + if await self.oauth.async_refresh_tokens(): + self.token_refresh_failure = False + await self._hass.async_add_executor_job(self.websocket.start_listener) + + return await self._hass.async_add_executor_job(func, *args) + + async def async_update(self) -> dict | None: + """Get the latest availability data.""" + data = await self.update_call( + self.fsr.get_availability, str(self._hass.config.time_zone) + ) + + if not data: + return None + + self.on_duty = bool(data.get("available")) + + _LOGGER.debug("Updated availability data: %s", data) + return data + + async def async_response_update(self) -> dict | None: + """Get the latest incident response data.""" + + if not self.incident_id: + return None + + _LOGGER.debug("Updating response data for incident id %s", self.incident_id) + + return await self.update_call(self.fsr.get_incident_response, self.incident_id) + + async def async_set_response(self, value) -> None: + """Set incident response status.""" + + if not self.incident_id: + return + + _LOGGER.debug( + "Setting incident response for incident id '%s' to state '%s'", + self.incident_id, + value, + ) + + await self.update_call(self.fsr.set_incident_response, self.incident_id, value) diff --git a/tests/components/fireservicerota/test_config_flow.py b/tests/components/fireservicerota/test_config_flow.py index 5555a8d649c..8d150034ec9 100644 --- a/tests/components/fireservicerota/test_config_flow.py +++ b/tests/components/fireservicerota/test_config_flow.py @@ -66,7 +66,7 @@ async def test_invalid_credentials(hass: HomeAssistant) -> None: """Test that invalid credentials throws an error.""" with patch( - "homeassistant.components.fireservicerota.FireServiceRota.request_tokens", + "homeassistant.components.fireservicerota.coordinator.FireServiceRota.request_tokens", side_effect=InvalidAuthError, ): result = await hass.config_entries.flow.async_init(