mirror of https://github.com/home-assistant/core
93 lines
2.3 KiB
Python
93 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()
|