* Refactor rate limit helper to track time in seconds Currently we created datetime and timedelta objects to enforce the rate limit. When the rate limit was being hit hard, this got expensive. We now use floats everywhere instead as they are much cheaper which is important when we are running up against a rate limit, which is by definition a hot path The rate limit helper is currently only used for templates and we do not have any code in the code base that directly passes in a rate limit so the impact to custom components is expected to be negligible if any * misesd two
92 lines
2.3 KiB
Python
92 lines
2.3 KiB
Python
"""Tests for ratelimit."""
|
|
|
|
import asyncio
|
|
import time
|
|
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers import ratelimit
|
|
|
|
|
|
async def test_hit(hass: HomeAssistant) -> None:
|
|
"""Test hitting the rate limit."""
|
|
|
|
refresh_called = False
|
|
|
|
@callback
|
|
def _refresh():
|
|
nonlocal refresh_called
|
|
refresh_called = True
|
|
|
|
rate_limiter = ratelimit.KeyedRateLimit(hass)
|
|
rate_limiter.async_triggered("key1", time.time())
|
|
|
|
assert (
|
|
rate_limiter.async_schedule_action("key1", 0.001, time.time(), _refresh)
|
|
is not None
|
|
)
|
|
|
|
assert not refresh_called
|
|
|
|
assert rate_limiter.async_has_timer("key1")
|
|
|
|
await asyncio.sleep(0.002)
|
|
assert refresh_called
|
|
|
|
assert (
|
|
rate_limiter.async_schedule_action("key2", 0.001, time.time(), _refresh) is None
|
|
)
|
|
rate_limiter.async_remove()
|
|
|
|
|
|
async def test_miss(hass: HomeAssistant) -> None:
|
|
"""Test missing the rate limit."""
|
|
|
|
refresh_called = False
|
|
|
|
@callback
|
|
def _refresh():
|
|
nonlocal refresh_called
|
|
refresh_called = True
|
|
|
|
rate_limiter = ratelimit.KeyedRateLimit(hass)
|
|
assert (
|
|
rate_limiter.async_schedule_action("key1", 0.1, time.time(), _refresh) is None
|
|
)
|
|
assert not refresh_called
|
|
assert not rate_limiter.async_has_timer("key1")
|
|
|
|
assert (
|
|
rate_limiter.async_schedule_action("key1", 0.1, time.time(), _refresh) is None
|
|
)
|
|
assert not refresh_called
|
|
assert not rate_limiter.async_has_timer("key1")
|
|
rate_limiter.async_remove()
|
|
|
|
|
|
async def test_no_limit(hass: HomeAssistant) -> None:
|
|
"""Test async_schedule_action always return None when there is no rate limit."""
|
|
|
|
refresh_called = False
|
|
|
|
@callback
|
|
def _refresh():
|
|
nonlocal refresh_called
|
|
refresh_called = True
|
|
|
|
rate_limiter = ratelimit.KeyedRateLimit(hass)
|
|
rate_limiter.async_triggered("key1", time.time())
|
|
|
|
assert (
|
|
rate_limiter.async_schedule_action("key1", None, time.time(), _refresh) is None
|
|
)
|
|
assert not refresh_called
|
|
assert not rate_limiter.async_has_timer("key1")
|
|
|
|
rate_limiter.async_triggered("key1", time.time())
|
|
|
|
assert (
|
|
rate_limiter.async_schedule_action("key1", None, time.time(), _refresh) is None
|
|
)
|
|
assert not refresh_called
|
|
assert not rate_limiter.async_has_timer("key1")
|
|
rate_limiter.async_remove()
|