diff --git a/homeassistant/components/nextbus/__init__.py b/homeassistant/components/nextbus/__init__.py index e8c0bc224febc3..817990620fe75b 100644 --- a/homeassistant/components/nextbus/__init__.py +++ b/homeassistant/components/nextbus/__init__.py @@ -13,15 +13,19 @@ 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 781742e4c08a06..dcaafa9573b750 100644 --- a/homeassistant/components/nextbus/coordinator.py +++ b/homeassistant/components/nextbus/coordinator.py @@ -48,27 +48,63 @@ def has_routes(self) -> bool: """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 + + 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) + ) - if prediction_results: - predictions[route_stop] = prediction_results[0] self._predictions = predictions return predictions diff --git a/homeassistant/components/nextbus/sensor.py b/homeassistant/components/nextbus/sensor.py index 8ef5323858f985..554814fe2db11c 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 231faccf907388..03e62a811f4198 100644 --- a/tests/components/nextbus/conftest.py +++ b/tests/components/nextbus/conftest.py @@ -41,7 +41,7 @@ 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 dd0346c3e7a9b8..8b62ed453b263e 100644 --- a/tests/components/nextbus/test_sensor.py +++ b/tests/components/nextbus/test_sensor.py @@ -5,6 +5,7 @@ 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.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_2 = { + DOMAIN: { + CONF_AGENCY: VALID_AGENCY, + CONF_ROUTE: ROUTE_2, + CONF_STOP: VALID_STOP, + } +} + BASIC_RESULTS = [ { "route": { @@ -60,7 +74,20 @@ {"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 @@ "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"