Throttle Nextbus if we are reaching the rate limit (#146064)

Co-authored-by: Josef Zweck <josef@zweck.dev>
Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/nextbus/coordinator.py

@@ -1,8 +1,8 @@
 """NextBus data update coordinator."""
 
-from datetime import timedelta
+from datetime import datetime, timedelta
 import logging
-from typing import Any
+from typing import Any, override
 
 from py_nextbus import NextBusClient
 from py_nextbus.client import NextBusFormatError, NextBusHTTPError
@@ -15,8 +15,14 @@ from .util import RouteStop
 
 _LOGGER = logging.getLogger(__name__)
 
+# At what percentage of the request limit should the coordinator pause making requests
+UPDATE_INTERVAL_SECONDS = 30
+THROTTLE_PRECENTAGE = 80
 
-class NextBusDataUpdateCoordinator(DataUpdateCoordinator):
+
+class NextBusDataUpdateCoordinator(
+    DataUpdateCoordinator[dict[RouteStop, dict[str, Any]]]
+):
     """Class to manage fetching NextBus data."""
 
     def __init__(self, hass: HomeAssistant, agency: str) -> None:
@@ -26,7 +32,7 @@ class NextBusDataUpdateCoordinator(DataUpdateCoordinator):
             _LOGGER,
             config_entry=None,  # It is shared between multiple entries
             name=DOMAIN,
-            update_interval=timedelta(seconds=30),
+            update_interval=timedelta(seconds=UPDATE_INTERVAL_SECONDS),
         )
         self.client = NextBusClient(agency_id=agency)
         self._agency = agency
@@ -49,9 +55,26 @@ class NextBusDataUpdateCoordinator(DataUpdateCoordinator):
         """Check if this coordinator is tracking any routes."""
         return len(self._route_stops) > 0
 
-    async def _async_update_data(self) -> dict[str, Any]:
+    @override
+    async def _async_update_data(self) -> dict[RouteStop, dict[str, Any]]:
         """Fetch data from NextBus."""
 
+        if (
+            # If we have predictions, check the rate limit
+            self._predictions
+            # If are over our rate limit percentage, we should throttle
+            and self.client.rate_limit_percent >= THROTTLE_PRECENTAGE
+            # But only if we have a reset time to unthrottle
+            and self.client.rate_limit_reset is not None
+            # Unless we are after the reset time
+            and datetime.now() < self.client.rate_limit_reset
+        ):
+            self.logger.debug(
+                "Rate limit threshold reached. Skipping updates for. Routes: %s",
+                str(self._route_stops),
+            )
+            return self._predictions
+
         _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)
@@ -60,7 +83,7 @@ class NextBusDataUpdateCoordinator(DataUpdateCoordinator):
             "Updating data from API. Routes: %s", str(_stops_to_route_stops)
         )
 
-        def _update_data() -> dict:
+        def _update_data() -> dict[RouteStop, dict[str, Any]]:
             """Fetch data from NextBus."""
             self.logger.debug("Updating data from API (executor)")
             predictions: dict[RouteStop, dict[str, Any]] = {}

tests/components/nextbus/conftest.py

@@ -137,6 +137,13 @@ def mock_nextbus_lists(
 def mock_nextbus() -> Generator[MagicMock]:
     """Create a mock py_nextbus module."""
     with patch("homeassistant.components.nextbus.coordinator.NextBusClient") as client:
+        instance = client.return_value
+
+        # Set some mocked rate limit values
+        instance.rate_limit = 450
+        instance.rate_limit_remaining = 225
+        instance.rate_limit_percent = 50.0
+
         yield client
 
 

tests/components/nextbus/test_sensor.py

@@ -1,6 +1,7 @@
 """The tests for the nexbus sensor component."""
 
 from copy import deepcopy
+from datetime import datetime, timedelta
 from unittest.mock import MagicMock
 from urllib.error import HTTPError
 
@@ -122,6 +123,57 @@ async def test_verify_no_upcoming(
     assert state.state == "unknown"
 
 
+async def test_verify_throttle(
+    hass: HomeAssistant,
+    mock_nextbus: MagicMock,
+    mock_nextbus_lists: MagicMock,
+    mock_nextbus_predictions: MagicMock,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Verify that the sensor coordinator is throttled correctly."""
+
+    # Set rate limit past threshold, should be ignored for first request
+    mock_client = mock_nextbus.return_value
+    mock_client.rate_limit_percent = 99.0
+    mock_client.rate_limit_reset = datetime.now() + timedelta(seconds=30)
+
+    # Do a request with the initial config and get predictions
+    await assert_setup_sensor(hass, CONFIG_BASIC)
+
+    # Validate the predictions are present
+    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"
+
+    # Update the predictions mock to return a different result
+    mock_nextbus_predictions.return_value = NO_UPCOMING
+
+    # Move time forward and bump the rate limit reset time
+    mock_client.rate_limit_reset = freezer.tick(31) + timedelta(seconds=30)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    # Verify that the sensor state is unchanged
+    state = hass.states.get(SENSOR_ID)
+    assert state is not None
+    assert state.state == "2019-03-28T21:09:31+00:00"
+
+    # Move time forward past the rate limit reset time
+    freezer.tick(31)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    # Verify that the sensor state is updated with the new predictions
+    state = hass.states.get(SENSOR_ID)
+    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,