From f448f488ba9ce1991b94bf52500b5e8f5a96f02a Mon Sep 17 00:00:00 2001 From: Ian Date: Tue, 10 Jun 2025 06:03:20 -0700 Subject: [PATCH] Throttle Nextbus if we are reaching the rate limit (#146064) Co-authored-by: Josef Zweck Co-authored-by: Robert Resch --- .../components/nextbus/coordinator.py | 35 ++++++++++--- tests/components/nextbus/conftest.py | 7 +++ tests/components/nextbus/test_sensor.py | 52 +++++++++++++++++++ 3 files changed, 88 insertions(+), 6 deletions(-) diff --git a/homeassistant/components/nextbus/coordinator.py b/homeassistant/components/nextbus/coordinator.py index 617669adf2f..e8d7ab06915 100644 --- a/homeassistant/components/nextbus/coordinator.py +++ b/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]] = {} diff --git a/tests/components/nextbus/conftest.py b/tests/components/nextbus/conftest.py index 3f687989313..9891f6ffa49 100644 --- a/tests/components/nextbus/conftest.py +++ b/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 diff --git a/tests/components/nextbus/test_sensor.py b/tests/components/nextbus/test_sensor.py index 04140a17c4f..eacab5cd5c4 100644 --- a/tests/components/nextbus/test_sensor.py +++ b/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,