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
home-assistant · Sep 20, 2024 · dccdb71 · dccdb71
1 parent df0195b
commit dccdb71
Show file tree

Hide file tree

Showing 5 changed files with 212 additions and 60 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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