hass-core/tests/helpers/test_ratelimit.py
J. Nick Koston b574220247
Refactor rate limit helper to track time in seconds (#113898)
* 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
2024-03-20 19:49:37 -04:00

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()