From dccdb71b2d07c515321f29ec5fbc395586ec131e Mon Sep 17 00:00:00 2001 From: Ian Date: Fri, 20 Sep 2024 01:18:13 -0700 Subject: [PATCH] Make NextBus coordinator more resilient and efficient (#126161) * Make NextBus coordinator more resilient and efficient Resolves issues where one request failing will prevent all agency predictions to fail. This also removes redundant requests for predictions that share the same stop. * Add unload entry test * Prevent shutdown if the coordinator is still needed --- homeassistant/components/nextbus/__init__.py | 25 +++-- .../components/nextbus/coordinator.py | 54 +++++++-- homeassistant/components/nextbus/sensor.py | 4 +- tests/components/nextbus/conftest.py | 86 +++++++++------ tests/components/nextbus/test_sensor.py | 103 +++++++++++++++++- 5 files changed, 212 insertions(+), 60 deletions(-) diff --git a/homeassistant/components/nextbus/__init__.py b/homeassistant/components/nextbus/__init__.py index e8c0bc224fe..817990620fe 100644 --- a/homeassistant/components/nextbus/__init__.py +++ b/homeassistant/components/nextbus/__init__.py @@ -13,15 +13,19 @@ PLATFORMS = [Platform.SENSOR] async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up platforms for NextBus.""" entry_agency = entry.data[CONF_AGENCY] + entry_stop = entry.data[CONF_STOP] + coordinator_key = f"{entry_agency}-{entry_stop}" - coordinator: NextBusDataUpdateCoordinator = hass.data.setdefault(DOMAIN, {}).get( - entry_agency + coordinator: NextBusDataUpdateCoordinator | None = hass.data.setdefault( + DOMAIN, {} + ).get( + coordinator_key, ) if coordinator is None: coordinator = NextBusDataUpdateCoordinator(hass, entry_agency) - hass.data[DOMAIN][entry_agency] = coordinator + hass.data[DOMAIN][coordinator_key] = coordinator - coordinator.add_stop_route(entry.data[CONF_STOP], entry.data[CONF_ROUTE]) + coordinator.add_stop_route(entry_stop, entry.data[CONF_ROUTE]) await coordinator.async_config_entry_first_refresh() @@ -33,11 +37,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" if await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - entry_agency = entry.data.get(CONF_AGENCY) - coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][entry_agency] - coordinator.remove_stop_route(entry.data[CONF_STOP], entry.data[CONF_ROUTE]) + entry_agency = entry.data[CONF_AGENCY] + entry_stop = entry.data[CONF_STOP] + coordinator_key = f"{entry_agency}-{entry_stop}" + + coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][coordinator_key] + coordinator.remove_stop_route(entry_stop, entry.data[CONF_ROUTE]) + if not coordinator.has_routes(): - hass.data[DOMAIN].pop(entry_agency) + await coordinator.async_shutdown() + hass.data[DOMAIN].pop(coordinator_key) return True diff --git a/homeassistant/components/nextbus/coordinator.py b/homeassistant/components/nextbus/coordinator.py index 781742e4c08..dcaafa9573b 100644 --- a/homeassistant/components/nextbus/coordinator.py +++ b/homeassistant/components/nextbus/coordinator.py @@ -48,27 +48,63 @@ class NextBusDataUpdateCoordinator(DataUpdateCoordinator): """Check if this coordinator is tracking any routes.""" return len(self._route_stops) > 0 + async def async_shutdown(self) -> None: + """If there are no more routes, cancel any scheduled call, and ignore new runs.""" + if self.has_routes(): + return + + await super().async_shutdown() + async def _async_update_data(self) -> dict[str, Any]: """Fetch data from NextBus.""" - _route_stops = set(self._route_stops) - self.logger.debug("Updating data from API. Routes: %s", str(_route_stops)) + _stops_to_route_stops: dict[str, set[RouteStop]] = {} + for route_stop in self._route_stops: + _stops_to_route_stops.setdefault(route_stop.stop_id, set()).add(route_stop) + + self.logger.debug( + "Updating data from API. Routes: %s", str(_stops_to_route_stops) + ) def _update_data() -> dict: """Fetch data from NextBus.""" self.logger.debug("Updating data from API (executor)") predictions: dict[RouteStop, dict[str, Any]] = {} - for route_stop in _route_stops: - prediction_results: list[dict[str, Any]] = [] + + for stop_id, route_stops in _stops_to_route_stops.items(): + self.logger.debug("Updating data from API (executor) %s", stop_id) try: - prediction_results = self.client.predictions_for_stop( - route_stop.stop_id, route_stop.route_id + prediction_results = self.client.predictions_for_stop(stop_id) + except NextBusHTTPError as ex: + self.logger.error( + "Error updating %s (executor): %s %s", + str(stop_id), + ex, + getattr(ex, "response", None), ) - except (NextBusHTTPError, NextBusFormatError) as ex: + raise UpdateFailed("Failed updating nextbus data", ex) from ex + except NextBusFormatError as ex: raise UpdateFailed("Failed updating nextbus data", ex) from ex - if prediction_results: - predictions[route_stop] = prediction_results[0] + self.logger.debug( + "Prediction results for %s (executor): %s", + str(stop_id), + str(prediction_results), + ) + + for route_stop in route_stops: + for prediction_result in prediction_results: + if ( + prediction_result["stop"]["id"] == route_stop.stop_id + and prediction_result["route"]["id"] == route_stop.route_id + ): + predictions[route_stop] = prediction_result + break + else: + self.logger.warning( + "Prediction not found for %s (executor)", str(route_stop) + ) + self._predictions = predictions return predictions diff --git a/homeassistant/components/nextbus/sensor.py b/homeassistant/components/nextbus/sensor.py index 8ef5323858f..554814fe2db 100644 --- a/homeassistant/components/nextbus/sensor.py +++ b/homeassistant/components/nextbus/sensor.py @@ -28,8 +28,10 @@ async def async_setup_entry( """Load values from configuration and initialize the platform.""" _LOGGER.debug(config.data) entry_agency = config.data[CONF_AGENCY] + entry_stop = config.data[CONF_STOP] + coordinator_key = f"{entry_agency}-{entry_stop}" - coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN].get(entry_agency) + coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN].get(coordinator_key) async_add_entities( ( diff --git a/tests/components/nextbus/conftest.py b/tests/components/nextbus/conftest.py index 231faccf907..03e62a811f4 100644 --- a/tests/components/nextbus/conftest.py +++ b/tests/components/nextbus/conftest.py @@ -41,7 +41,7 @@ import pytest def route_config_direction(request: pytest.FixtureRequest) -> Any: """Generate alternative directions values. - When only on edirection is returned, it is not returned as a list, but instead an object. + When only one direction is returned, it is not returned as a list, but instead an object. """ return request.param @@ -75,42 +75,56 @@ def mock_nextbus_lists( "hidden": False, "timestamp": "2024-06-23T03:06:58Z", }, + { + "id": "G", + "rev": 1057, + "title": "F Market & Wharves", + "description": "7am-10pm daily", + "color": "", + "textColor": "", + "hidden": False, + "timestamp": "2024-06-23T03:06:58Z", + }, ] - instance.route_details.return_value = { - "id": "F", - "rev": 1057, - "title": "F Market & Wharves", - "description": "7am-10pm daily", - "color": "", - "textColor": "", - "hidden": False, - "boundingBox": {}, - "stops": [ - { - "id": "5184", - "lat": 37.8071299, - "lon": -122.41732, - "name": "Jones St & Beach St", - "code": "15184", - "hidden": False, - "showDestinationSelector": True, - "directions": ["F_0_var1", "F_0_var0"], - }, - { - "id": "5651", - "lat": 37.8071299, - "lon": -122.41732, - "name": "Jones St & Beach St", - "code": "15651", - "hidden": False, - "showDestinationSelector": True, - "directions": ["F_0_var1", "F_0_var0"], - }, - ], - "directions": route_config_direction, - "paths": [], - "timestamp": "2024-06-23T03:06:58Z", - } + def route_details_side_effect(agency: str, route: str) -> dict: + route = route.upper() + return { + "id": route, + "rev": 1057, + "title": f"{route} Market & Wharves", + "description": "7am-10pm daily", + "color": "", + "textColor": "", + "hidden": False, + "boundingBox": {}, + "stops": [ + { + "id": "5184", + "lat": 37.8071299, + "lon": -122.41732, + "name": "Jones St & Beach St", + "code": "15184", + "hidden": False, + "showDestinationSelector": True, + "directions": ["F_0_var1", "F_0_var0"], + }, + { + "id": "5651", + "lat": 37.8071299, + "lon": -122.41732, + "name": "Jones St & Beach St", + "code": "15651", + "hidden": False, + "showDestinationSelector": True, + "directions": ["F_0_var1", "F_0_var0"], + }, + ], + "directions": route_config_direction, + "paths": [], + "timestamp": "2024-06-23T03:06:58Z", + } + + instance.route_details.side_effect = route_details_side_effect return instance diff --git a/tests/components/nextbus/test_sensor.py b/tests/components/nextbus/test_sensor.py index dd0346c3e7a..8b62ed453b2 100644 --- a/tests/components/nextbus/test_sensor.py +++ b/tests/components/nextbus/test_sensor.py @@ -5,6 +5,7 @@ from copy import deepcopy from unittest.mock import MagicMock, patch from urllib.error import HTTPError +from freezegun.api import FrozenDateTimeFactory from py_nextbus.client import NextBusFormatError, NextBusHTTPError import pytest @@ -16,16 +17,21 @@ from homeassistant.const import CONF_NAME, CONF_STOP from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import UpdateFailed -from tests.common import MockConfigEntry +from tests.common import MockConfigEntry, async_fire_time_changed VALID_AGENCY = "sfmta-cis" VALID_ROUTE = "F" VALID_STOP = "5184" +VALID_COORDINATOR_KEY = f"{VALID_AGENCY}-{VALID_STOP}" VALID_AGENCY_TITLE = "San Francisco Muni" VALID_ROUTE_TITLE = "F-Market & Wharves" VALID_STOP_TITLE = "Market St & 7th St" SENSOR_ID = "sensor.san_francisco_muni_f_market_wharves_market_st_7th_st" +ROUTE_2 = "G" +ROUTE_TITLE_2 = "G-Market & Wharves" +SENSOR_ID_2 = "sensor.san_francisco_muni_g_market_wharves_market_st_7th_st" + PLATFORM_CONFIG = { sensor.DOMAIN: { "platform": DOMAIN, @@ -44,6 +50,14 @@ CONFIG_BASIC = { } } +CONFIG_BASIC_2 = { + DOMAIN: { + CONF_AGENCY: VALID_AGENCY, + CONF_ROUTE: ROUTE_2, + CONF_STOP: VALID_STOP, + } +} + BASIC_RESULTS = [ { "route": { @@ -60,7 +74,20 @@ BASIC_RESULTS = [ {"minutes": 3, "timestamp": 1553807373000}, {"minutes": 10, "timestamp": 1553807380000}, ], - } + }, + { + "route": { + "title": ROUTE_TITLE_2, + "id": ROUTE_2, + }, + "stop": { + "name": VALID_STOP_TITLE, + "id": VALID_STOP, + }, + "values": [ + {"minutes": 90, "timestamp": 1553807379000}, + ], + }, ] NO_UPCOMING = [ @@ -74,7 +101,18 @@ NO_UPCOMING = [ "id": VALID_STOP, }, "values": [], - } + }, + { + "route": { + "title": ROUTE_TITLE_2, + "id": ROUTE_2, + }, + "stop": { + "name": VALID_STOP_TITLE, + "id": VALID_STOP, + }, + "values": [], + }, ] @@ -100,13 +138,15 @@ async def assert_setup_sensor( hass: HomeAssistant, config: dict[str, dict[str, str]], expected_state=ConfigEntryState.LOADED, + route_title: str = VALID_ROUTE_TITLE, ) -> MockConfigEntry: """Set up the sensor and assert it's been created.""" + unique_id = f"{config[DOMAIN][CONF_AGENCY]}_{config[DOMAIN][CONF_ROUTE]}_{config[DOMAIN][CONF_STOP]}" config_entry = MockConfigEntry( domain=DOMAIN, data=config[DOMAIN], - title=f"{VALID_AGENCY_TITLE} {VALID_ROUTE_TITLE} {VALID_STOP_TITLE}", - unique_id=f"{VALID_AGENCY}_{VALID_ROUTE}_{VALID_STOP}", + title=f"{VALID_AGENCY_TITLE} {route_title} {VALID_STOP_TITLE}", + unique_id=unique_id, ) config_entry.add_to_hass(hass) @@ -153,7 +193,7 @@ async def test_prediction_exceptions( ) -> None: """Test that some coodinator exceptions raise UpdateFailed exceptions.""" await assert_setup_sensor(hass, CONFIG_BASIC) - coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][VALID_AGENCY] + coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][VALID_COORDINATOR_KEY] mock_nextbus_predictions.side_effect = client_exception with pytest.raises(UpdateFailed): await coordinator._async_update_data() @@ -205,3 +245,54 @@ async def test_verify_no_upcoming( assert state is not None assert state.attributes["upcoming"] == "No upcoming predictions" assert state.state == "unknown" + + +async def test_unload_entry( + hass: HomeAssistant, + mock_nextbus: MagicMock, + mock_nextbus_lists: MagicMock, + mock_nextbus_predictions: MagicMock, + freezer: FrozenDateTimeFactory, +) -> None: + """Test that the sensor can be unloaded.""" + config_entry1 = await assert_setup_sensor(hass, CONFIG_BASIC) + await assert_setup_sensor(hass, CONFIG_BASIC_2, route_title=ROUTE_TITLE_2) + + # Verify the first sensor + state = hass.states.get(SENSOR_ID) + assert state is not None + assert state.state == "2019-03-28T21:09:31+00:00" + assert state.attributes["agency"] == VALID_AGENCY + assert state.attributes["route"] == VALID_ROUTE_TITLE + assert state.attributes["stop"] == VALID_STOP_TITLE + assert state.attributes["upcoming"] == "1, 2, 3, 10" + + # Verify the second sensor + state = hass.states.get(SENSOR_ID_2) + assert state is not None + assert state.state == "2019-03-28T21:09:39+00:00" + assert state.attributes["agency"] == VALID_AGENCY + assert state.attributes["route"] == ROUTE_TITLE_2 + assert state.attributes["stop"] == VALID_STOP_TITLE + assert state.attributes["upcoming"] == "90" + + # Update mock to return new predictions + new_predictions = deepcopy(BASIC_RESULTS) + new_predictions[1]["values"] = [{"minutes": 5, "timestamp": 1553807375000}] + mock_nextbus_predictions.return_value = new_predictions + + # Unload config entry 1 + await hass.config_entries.async_unload(config_entry1.entry_id) + await hass.async_block_till_done() + assert config_entry1.state is ConfigEntryState.NOT_LOADED + + # Skip ahead in time + freezer.tick(120) + async_fire_time_changed(hass) + await hass.async_block_till_done(wait_background_tasks=True) + + # Check update for new predictions + state = hass.states.get(SENSOR_ID_2) + assert state is not None + assert state.attributes["upcoming"] == "5" + assert state.state == "2019-03-28T21:09:35+00:00"