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

@@ -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(
+    coordinator: NextBusDataUpdateCoordinator | None = hass.data.setdefault(
-        entry_agency
+        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)
+        entry_agency = entry.data[CONF_AGENCY]
-        coordinator: NextBusDataUpdateCoordinator = hass.data[DOMAIN][entry_agency]
+        entry_stop = entry.data[CONF_STOP]
-        coordinator.remove_stop_route(entry.data[CONF_STOP], entry.data[CONF_ROUTE])
+        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
 

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)
+        _stops_to_route_stops: dict[str, set[RouteStop]] = {}
-        self.logger.debug("Updating data from API. Routes: %s", str(_route_stops))
+        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(
+                    prediction_results = self.client.predictions_for_stop(stop_id)
-                        route_stop.stop_id, route_stop.route_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:
+                self.logger.debug(
-                    predictions[route_stop] = prediction_results[0]
+                    "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

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(
         (

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 = {
+    def route_details_side_effect(agency: str, route: str) -> dict:
-        "id": "F",
+        route = route.upper()
-        "rev": 1057,
+        return {
-        "title": "F Market & Wharves",
+            "id": route,
-        "description": "7am-10pm daily",
+            "rev": 1057,
-        "color": "",
+            "title": f"{route} Market & Wharves",
-        "textColor": "",
+            "description": "7am-10pm daily",
-        "hidden": False,
+            "color": "",
-        "boundingBox": {},
+            "textColor": "",
-        "stops": [
+            "hidden": False,
-            {
+            "boundingBox": {},
-                "id": "5184",
+            "stops": [
-                "lat": 37.8071299,
+                {
-                "lon": -122.41732,
+                    "id": "5184",
-                "name": "Jones St & Beach St",
+                    "lat": 37.8071299,
-                "code": "15184",
+                    "lon": -122.41732,
-                "hidden": False,
+                    "name": "Jones St & Beach St",
-                "showDestinationSelector": True,
+                    "code": "15184",
-                "directions": ["F_0_var1", "F_0_var0"],
+                    "hidden": False,
-            },
+                    "showDestinationSelector": True,
-            {
+                    "directions": ["F_0_var1", "F_0_var0"],
-                "id": "5651",
+                },
-                "lat": 37.8071299,
+                {
-                "lon": -122.41732,
+                    "id": "5651",
-                "name": "Jones St & Beach St",
+                    "lat": 37.8071299,
-                "code": "15651",
+                    "lon": -122.41732,
-                "hidden": False,
+                    "name": "Jones St & Beach St",
-                "showDestinationSelector": True,
+                    "code": "15651",
-                "directions": ["F_0_var1", "F_0_var0"],
+                    "hidden": False,
-            },
+                    "showDestinationSelector": True,
-        ],
+                    "directions": ["F_0_var1", "F_0_var0"],
-        "directions": route_config_direction,
+                },
-        "paths": [],
+            ],
-        "timestamp": "2024-06-23T03:06:58Z",
+            "directions": route_config_direction,
-    }
+            "paths": [],
+            "timestamp": "2024-06-23T03:06:58Z",
+        }
+
+    instance.route_details.side_effect = route_details_side_effect
 
     return instance

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}",
+        title=f"{VALID_AGENCY_TITLE} {route_title} {VALID_STOP_TITLE}",
-        unique_id=f"{VALID_AGENCY}_{VALID_ROUTE}_{VALID_STOP}",
+        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"