8000 NextBus: Self-throttle using client rate-limit data · home-assistant/core@8cc64b0 · GitHub
[go: up one dir, main page]
Skip to content

Commit 8cc64b0

Browse files
ViViDboarderViViDboarder
committed
NextBus: Self-throttle using client rate-limit data
This should help avoid rate limit exceeded errors
1 parent a1c6441 commit 8cc64b0

File tree

3 files changed

+86
-7
lines changed

3 files changed

+86
-7
lines changed

‎homeassistant/components/nextbus/coordinator.py

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
"""NextBus data update coordinator."""
22

3-
from datetime import timedelta
3+
from datetime import datetime, timedelta
44
import logging
55
from typing import Any
66

@@ -15,8 +15,14 @@
1515

1616
_LOGGER = logging.getLogger(__name__)
1717

18+
# At what percentage of the request limit should the coordinator pause making requests
19+
UPDATE_INTERVAL_SECONDS = 30
20+
THROTTLE_PRECENTAGE = 80
1821

19-
class NextBusDataUpdateCoordinator(DataUpdateCoordinator):
22+
23+
class NextBusDataUpdateCoordinator(
24+
DataUpdateCoordinator[dict[RouteStop, dict[str, Any]]]
25+
):
2026
"""Class to manage fetching NextBus data."""
2127

2228
def __init__(self, hass: HomeAssistant, agency: str) -> None:
@@ -26,10 +32,10 @@ def __init__(self, hass: HomeAssistant, agency: str) -> None:
2632
_LOGGER,
2733
config_entry=None, # It is shared between multiple entries
2834
name=DOMAIN,
29-
update_interval=timedelta(seconds=30),
35+
update_interval=timedelta(seconds=UPDATE_INTERVAL_SECONDS),
3036
)
31-
self.client = NextBusClient(agency_id=agency)
32-
self._agency = agency
37+
self.client: NextBusClient = NextBusClient(agency_id=agency)
38+
self._agency: str = agency
3339
self._route_stops: set[RouteStop] = set()
3440
self._predictions: dict[RouteStop, dict[str, Any]] = {}
3541

@@ -49,9 +55,23 @@ def has_routes(self) -> bool:
4955
"""Check if this coordinator is tracking any routes."""
5056
return len(self._route_stops) > 0
5157

52-
async def _async_update_data(self) -> dict[str, Any]:
58+
async def _async_update_data(self) -> dict[RouteStop, dict[str, Any]]:
5359
"""Fetch data from NextBus."""
5460

61+
if (
62+
# If we have predictions, check the rate limit
63+
self._predictions
64+
# If are over our rate limit percentage, we should throttle
65+
and self.client.rate_limit_percent >= THROTTLE_PRECENTAGE
66+
# Unless we are after the reset time
67+
and datetime.now() < self.client.rate_limit_reset
68+
):
69+
self.logger.debug(
70+
"Rate limit threshold reached. Skipping updates for. Routes: %s",
71+
str(self._route_stops),
72+
)
73+
return self._predictions
74+
5575
_stops_to_route_stops: dict[str, set[RouteStop]] = {}
5676
for route_stop in self._route_stops:
5777
_stops_to_route_stops.setdefault(route_stop.stop_id, set()).add(route_stop)
@@ -60,7 +80,7 @@ async def _async_update_data(self) -> dict[str, Any]:
6080
"Updating data from API. Routes: %s", str(_stops_to_route_stops)
6181
)
6282

63-
def _update_data() -> dict:
83+
def _update_data() -> dict[RouteStop, dict[str, Any]]:
6484
"""Fetch data from NextBus."""
6585
self.logger.debug("Updating data from API (executor)")
6686
predictions: dict[RouteStop, dict[str, Any]] = {}

‎tests/components/nextbus/conftest.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,13 @@ def route_details_side_effect(agency: str, route: str) -> dict:
137137
def mock_nextbus() -> Generator[MagicMock]:
138138
"""Create a mock py_nextbus module."""
139139
with patch("homeassistant.components.nextbus.coordinator.NextBusClient") as client:
140+
instance = client.return_value
141+
142+
# Set some mocked rate limit values
143+
instance.rate_limit = 450
144+
instance.rate_limit_remaining = 225
145+
instance.rate_limit_percent = 50.0
146+
140147
yield client
141148

142149

‎tests/components/nextbus/test_sensor.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""The tests for the nexbus sensor component."""
22

33
from copy import deepcopy
4+
from datetime import datetime, timedelta
45
from unittest.mock import MagicMock
56
from urllib.error import HTTPError
67

@@ -122,6 +123,57 @@ async def test_verify_no_upcoming(
122123
assert state.state == "unknown"
123124

124125

126+
async def test_verify_throttle(
127+
hass: HomeAssistant,
128+
mock_nextbus: MagicMock,
129+
mock_nextbus_lists: MagicMock,
130+
mock_nextbus_predictions: MagicMock,
131+
freezer: FrozenDateTimeFactory,
132+
) -> None:
133+
"""Verify that the sensor coordinator is throttled correctly."""
134+
135+
# Set rate limit past threshold, should be ignored for first request
136+
mock_client = mock_nextbus.return_value
137+
mock_client.rate_limit_percent = 99.0
138+
mock_client.rate_limit_reset = datetime.now() + timedelta(seconds=30)
139+
140+
# Do a request with the initial config and get predictions
141+
await assert_setup_sensor(hass, CONFIG_BASIC)
142+
143+
# Validate the predictions are present
144+
state = hass.states.get(SENSOR_ID)
145+
assert state is not None
146+
assert state.state == "2019-03-28T21:09:31+00:00"
147+
assert state.attributes["agency"] == VALID_AGENCY
148+
assert state.attributes["route"] == VALID_ROUTE_TITLE
149+
assert state.attributes["stop"] == VALID_STOP_TITLE
150+
assert state.attributes["upcoming"] == "1, 2, 3, 10"
151+
152+
# Update the predictions mock to return a different result
153+
mock_nextbus_predictions.return_value = NO_UPCOMING
154+
155+
# Move time forward and bump the rate limit reset time
156+
mock_client.rate_limit_reset = freezer.tick(31) + timedelta(seconds=30)
157+
async_fire_time_changed(hass)
158+
await hass.async_block_till_done(wait_background_tasks=True)
159+
160+
# Verify that the sensor state is unchanged
161+
state = hass.states.get(SENSOR_ID)
162+
assert state is not None
163+
assert state.state == "2019-03-28T21:09:31+00:00"
164+
165+
# Move time forward past the rate limit reset time
166+
freezer.tick(31)
167+
async_fire_time_changed(hass)
168+
await hass.async_block_till_done(wait_background_tasks=True)
169+
170+
# Verify that the sensor state is updated with the new predictions
171+
state = hass.states.get(SENSOR_ID)
172+
assert state is not None
173+
assert state.attributes["upcoming"] == "No upcoming predictions"
174+
assert state.state == "unknown"
175+
176+
125177
async def test_unload_entry(
126178
hass: HomeAssistant,
127179
mock_nextbus: MagicMock,

0 commit comments

Comments
 (0)
0