Implement template rate_limit directive (#40667)

pull/41013/head
J. Nick Koston 2020-10-01 14:39:44 -05:00 committed by GitHub
parent b3464c5087
commit b45215f1d2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 669 additions and 18 deletions

View File

@ -627,3 +627,7 @@ CLOUD_NEVER_EXPOSED_ENTITIES = ["group.all_locks"]
# The ID of the Home Assistant Cast App
CAST_APP_ID_HOMEASSISTANT = "B12CE3CA"
# The tracker error allow when converting
# loop time to human readable time
MAX_TIME_TRACKING_ERROR = 0.001

View File

@ -538,7 +538,7 @@ class Event:
event_type: str,
data: Optional[Dict[str, Any]] = None,
origin: EventOrigin = EventOrigin.local,
time_fired: Optional[int] = None,
time_fired: Optional[datetime.datetime] = None,
context: Optional[Context] = None,
) -> None:
"""Initialize a new event."""

View File

@ -27,6 +27,7 @@ from homeassistant.const import (
EVENT_STATE_CHANGED,
EVENT_TIME_CHANGED,
MATCH_ALL,
MAX_TIME_TRACKING_ERROR,
SUN_EVENT_SUNRISE,
SUN_EVENT_SUNSET,
)
@ -40,6 +41,7 @@ from homeassistant.core import (
)
from homeassistant.exceptions import TemplateError
from homeassistant.helpers.entity_registry import EVENT_ENTITY_REGISTRY_UPDATED
from homeassistant.helpers.ratelimit import KeyedRateLimit
from homeassistant.helpers.sun import get_astral_event_next
from homeassistant.helpers.template import RenderInfo, Template, result_as_boolean
from homeassistant.helpers.typing import TemplateVarsType
@ -47,8 +49,6 @@ from homeassistant.loader import bind_hass
from homeassistant.util import dt as dt_util
from homeassistant.util.async_ import run_callback_threadsafe
MAX_TIME_TRACKING_ERROR = 0.001
TRACK_STATE_CHANGE_CALLBACKS = "track_state_change_callbacks"
TRACK_STATE_CHANGE_LISTENER = "track_state_change_listener"
@ -88,10 +88,12 @@ class TrackTemplate:
The template is template to calculate.
The variables are variables to pass to the template.
The rate_limit is a rate limit on how often the template is re-rendered.
"""
template: Template
variables: TemplateVarsType
rate_limit: Optional[timedelta] = None
@dataclass
@ -724,6 +726,8 @@ class _TrackTemplateResultInfo:
self._track_templates = track_templates
self._last_result: Dict[Template, Union[str, TemplateError]] = {}
self._rate_limit = KeyedRateLimit(hass)
self._info: Dict[Template, RenderInfo] = {}
self._track_state_changes: Optional[_TrackStateChangeFiltered] = None
@ -763,6 +767,7 @@ class _TrackTemplateResultInfo:
"""Cancel the listener."""
assert self._track_state_changes
self._track_state_changes.async_remove()
self._rate_limit.async_remove()
@callback
def async_refresh(self) -> None:
@ -784,11 +789,23 @@ class _TrackTemplateResultInfo:
def _refresh(self, event: Optional[Event]) -> None:
updates = []
info_changed = False
now = dt_util.utcnow()
for track_template_ in self._track_templates:
template = track_template_.template
if event:
if not self._event_triggers_template(template, event):
if not self._rate_limit.async_has_timer(
template
) and not self._event_triggers_template(template, event):
continue
if self._rate_limit.async_schedule_action(
template,
self._info[template].rate_limit or track_template_.rate_limit,
now,
self._refresh,
event,
):
continue
_LOGGER.debug(
@ -797,6 +814,7 @@ class _TrackTemplateResultInfo:
event,
)
self._rate_limit.async_triggered(template, now)
self._info[template] = template.async_render_to_info(
track_template_.variables
)

View File

@ -0,0 +1,97 @@
"""Ratelimit helper."""
import asyncio
from datetime import datetime, timedelta
import logging
from typing import Any, Callable, Dict, Hashable, Optional
from homeassistant.const import MAX_TIME_TRACKING_ERROR
from homeassistant.core import HomeAssistant, callback
import homeassistant.util.dt as dt_util
_LOGGER = logging.getLogger(__name__)
class KeyedRateLimit:
"""Class to track rate limits."""
def __init__(
self,
hass: HomeAssistant,
):
"""Initialize ratelimit tracker."""
self.hass = hass
self._last_triggered: Dict[Hashable, datetime] = {}
self._rate_limit_timers: Dict[Hashable, asyncio.TimerHandle] = {}
@callback
def async_has_timer(self, key: Hashable) -> bool:
"""Check if a rate limit timer is running."""
return key in self._rate_limit_timers
@callback
def async_triggered(self, key: Hashable, now: Optional[datetime] = None) -> None:
"""Call when the action we are tracking was triggered."""
self.async_cancel_timer(key)
self._last_triggered[key] = now or dt_util.utcnow()
@callback
def async_cancel_timer(self, key: Hashable) -> None:
"""Cancel a rate limit time that will call the action."""
if not self.async_has_timer(key):
return
self._rate_limit_timers.pop(key).cancel()
@callback
def async_remove(self) -> None:
"""Remove all timers."""
for timer in self._rate_limit_timers.values():
timer.cancel()
self._rate_limit_timers.clear()
@callback
def async_schedule_action(
self,
key: Hashable,
rate_limit: Optional[timedelta],
now: datetime,
action: Callable,
*args: Any,
) -> Optional[datetime]:
"""Check rate limits and schedule an action if we hit the limit.
If the rate limit is hit:
Schedules the action for when the rate limit expires
if there are no pending timers. The action must
be called in async.
Returns the time the rate limit will expire
If the rate limit is not hit:
Return None
"""
if rate_limit is None or key not in self._last_triggered:
return None
next_call_time = self._last_triggered[key] + rate_limit
if next_call_time <= now:
self.async_cancel_timer(key)
return None
_LOGGER.debug(
"Reached rate limit of %s for %s and deferred action until %s",
rate_limit,
key,
next_call_time,
)
if key not in self._rate_limit_timers:
self._rate_limit_timers[key] = self.hass.loop.call_later(
(next_call_time - now).total_seconds() + MAX_TIME_TRACKING_ERROR,
action,
*args,
)
return next_call_time

View File

@ -72,6 +72,8 @@ _COLLECTABLE_STATE_ATTRIBUTES = {
"name",
}
DEFAULT_RATE_LIMIT = timedelta(seconds=1)
@bind_hass
def attach(hass: HomeAssistantType, obj: Any) -> None:
@ -198,10 +200,11 @@ class RenderInfo:
self.domains = set()
self.domains_lifecycle = set()
self.entities = set()
self.rate_limit = None
def __repr__(self) -> str:
"""Representation of RenderInfo."""
return f"<RenderInfo {self.template} all_states={self.all_states} all_states_lifecycle={self.all_states_lifecycle} domains={self.domains} domains_lifecycle={self.domains_lifecycle} entities={self.entities}>"
return f"<RenderInfo {self.template} all_states={self.all_states} all_states_lifecycle={self.all_states_lifecycle} domains={self.domains} domains_lifecycle={self.domains_lifecycle} entities={self.entities} rate_limit={self.rate_limit}>"
def _filter_domains_and_entities(self, entity_id: str) -> bool:
"""Template should re-render if the entity state changes when we match specific domains or entities."""
@ -221,16 +224,24 @@ class RenderInfo:
def _freeze_static(self) -> None:
self.is_static = True
self.entities = frozenset(self.entities)
self.domains = frozenset(self.domains)
self.domains_lifecycle = frozenset(self.domains_lifecycle)
self._freeze_sets()
self.all_states = False
def _freeze(self) -> None:
def _freeze_sets(self) -> None:
self.entities = frozenset(self.entities)
self.domains = frozenset(self.domains)
self.domains_lifecycle = frozenset(self.domains_lifecycle)
def _freeze(self) -> None:
self._freeze_sets()
if self.rate_limit is None and (
self.domains or self.domains_lifecycle or self.all_states or self.exception
):
# If the template accesses all states or an entire
# domain, and no rate limit is set, we use the default.
self.rate_limit = DEFAULT_RATE_LIMIT
if self.exception:
return
@ -478,6 +489,26 @@ class Template:
return 'Template("' + self.template + '")'
class RateLimit:
"""Class to control update rate limits."""
def __init__(self, hass: HomeAssistantType):
"""Initialize rate limit."""
self._hass = hass
def __call__(self, *args: Any, **kwargs: Any) -> str:
"""Handle a call to the class."""
render_info = self._hass.data.get(_RENDER_INFO)
if render_info is not None:
render_info.rate_limit = timedelta(*args, **kwargs)
return ""
def __repr__(self) -> str:
"""Representation of a RateLimit."""
return "<template RateLimit>"
class AllStates:
"""Class to expose all HA states as attributes."""
@ -1279,10 +1310,11 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
self.globals["is_state_attr"] = hassfunction(is_state_attr)
self.globals["state_attr"] = hassfunction(state_attr)
self.globals["states"] = AllStates(hass)
self.globals["rate_limit"] = RateLimit(hass)
def is_safe_callable(self, obj):
"""Test if callback is safe."""
return isinstance(obj, AllStates) or super().is_safe_callable(obj)
return isinstance(obj, (AllStates, RateLimit)) or super().is_safe_callable(obj)
def is_safe_attribute(self, obj, attr, value):
"""Test if attribute is safe."""

View File

@ -115,6 +115,7 @@ async def test_template_state_boolean(hass, calls):
async def test_template_position(hass, calls):
"""Test the position_template attribute."""
hass.states.async_set("cover.test", STATE_OPEN)
with assert_setup_component(1, "cover"):
assert await setup.async_setup_component(
hass,

View File

@ -1,5 +1,6 @@
"""The test for the Template sensor platform."""
from asyncio import Event
from datetime import timedelta
from unittest.mock import patch
from homeassistant.bootstrap import async_from_config_dict
@ -17,7 +18,11 @@ from homeassistant.helpers.template import Template
from homeassistant.setup import ATTR_COMPONENT, async_setup_component, setup_component
import homeassistant.util.dt as dt_util
from tests.common import assert_setup_component, get_test_home_assistant
from tests.common import (
assert_setup_component,
async_fire_time_changed,
get_test_home_assistant,
)
class TestTemplateSensor:
@ -900,8 +905,13 @@ async def test_self_referencing_entity_picture_loop(hass, caplog):
assert len(hass.states.async_all()) == 1
await hass.async_block_till_done()
await hass.async_block_till_done()
next_time = dt_util.utcnow() + timedelta(seconds=1.2)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert "Template loop detected" in caplog.text

View File

@ -927,7 +927,7 @@ async def test_track_template_result_complex(hass):
"""Test tracking template."""
specific_runs = []
template_complex_str = """
{{ rate_limit(seconds=0) }}
{% if states("sensor.domain") == "light" %}
{{ states.light | map(attribute='entity_id') | list }}
{% elif states("sensor.domain") == "lock" %}
@ -1162,6 +1162,8 @@ async def test_track_template_result_with_group(hass):
await hass.services.async_call("group", "reload")
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert specific_runs[-1] == str(100.1 + 200.2 + 0 + 800.8)
@ -1234,7 +1236,7 @@ async def test_track_template_result_iterator(hass):
[
TrackTemplate(
Template(
"""
"""{{ rate_limit(seconds=0) }}
{% for state in states.sensor %}
{% if state.state == 'on' %}
{{ state.entity_id }},
@ -1266,7 +1268,7 @@ async def test_track_template_result_iterator(hass):
[
TrackTemplate(
Template(
"""{{ states.sensor|selectattr("state","equalto","on")
"""{{ rate_limit(seconds=0) }}{{ states.sensor|selectattr("state","equalto","on")
|join(",", attribute="entity_id") }}""",
hass,
),
@ -1397,6 +1399,335 @@ async def test_static_string(hass):
assert refresh_runs == ["static"]
async def test_track_template_rate_limit(hass):
"""Test template rate limit."""
template_refresh = Template("{{ states | count }}", hass)
refresh_runs = []
@ha.callback
def refresh_listener(event, updates):
refresh_runs.append(updates.pop().result)
info = async_track_template_result(
hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=0.1))],
refresh_listener,
)
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert refresh_runs == ["0"]
hass.states.async_set("sensor.one", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0"]
info.async_refresh()
assert refresh_runs == ["0", "1"]
hass.states.async_set("sensor.two", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.three", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.four", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "4"]
hass.states.async_set("sensor.five", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "4"]
async def test_track_template_rate_limit_overridden(hass):
"""Test template rate limit can be overridden from the template."""
template_refresh = Template(
"{% set x = rate_limit(seconds=0.1) %}{{ states | count }}", hass
)
refresh_runs = []
@ha.callback
def refresh_listener(event, updates):
refresh_runs.append(updates.pop().result)
info = async_track_template_result(
hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=5))],
refresh_listener,
)
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert refresh_runs == ["0"]
hass.states.async_set("sensor.one", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0"]
info.async_refresh()
assert refresh_runs == ["0", "1"]
hass.states.async_set("sensor.two", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.three", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.four", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "4"]
hass.states.async_set("sensor.five", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "4"]
async def test_track_template_rate_limit_five(hass):
"""Test template rate limit of 5 seconds."""
template_refresh = Template("{{ states | count }}", hass)
refresh_runs = []
@ha.callback
def refresh_listener(event, updates):
refresh_runs.append(updates.pop().result)
info = async_track_template_result(
hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=5))],
refresh_listener,
)
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert refresh_runs == ["0"]
hass.states.async_set("sensor.one", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0"]
info.async_refresh()
assert refresh_runs == ["0", "1"]
hass.states.async_set("sensor.two", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1"]
hass.states.async_set("sensor.three", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1"]
async def test_track_template_rate_limit_changes(hass):
"""Test template rate limit can be changed."""
template_refresh = Template(
"""
{% if states.sensor.two.state == "any" %}
{% set x = rate_limit(seconds=5) %}
{% else %}
{% set x = rate_limit(seconds=0.1) %}
{% endif %}
{{ states | count }}
""",
hass,
)
refresh_runs = []
@ha.callback
def refresh_listener(event, updates):
refresh_runs.append(updates.pop().result)
info = async_track_template_result(
hass,
[TrackTemplate(template_refresh, None)],
refresh_listener,
)
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert refresh_runs == ["0"]
hass.states.async_set("sensor.one", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0"]
info.async_refresh()
assert refresh_runs == ["0", "1"]
hass.states.async_set("sensor.two", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 1)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.three", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.four", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.five", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
async def test_track_template_rate_limit_removed(hass):
"""Test template rate limit can be removed."""
template_refresh = Template(
"""
{% if states.sensor.two.state == "any" %}
{% set x = rate_limit(0) %}
{% else %}
{% set x = rate_limit(seconds=0.1) %}
{% endif %}
{{ states | count }}
""",
hass,
)
refresh_runs = []
@ha.callback
def refresh_listener(event, updates):
refresh_runs.append(updates.pop().result)
info = async_track_template_result(
hass,
[TrackTemplate(template_refresh, None)],
refresh_listener,
)
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert refresh_runs == ["0"]
hass.states.async_set("sensor.one", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0"]
info.async_refresh()
assert refresh_runs == ["0", "1"]
hass.states.async_set("sensor.two", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 1)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2"]
hass.states.async_set("sensor.three", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "3"]
hass.states.async_set("sensor.four", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "3", "4"]
hass.states.async_set("sensor.five", "any")
await hass.async_block_till_done()
assert refresh_runs == ["0", "1", "2", "3", "4", "5"]
async def test_track_two_templates_with_different_rate_limits(hass):
"""Test two templates with different rate limits."""
template_one = Template(
"{% set x = rate_limit(seconds=0.1) %}{{ states | count }}", hass
)
template_five = Template(
"{% set x = rate_limit(seconds=5) %}{{ states | count }}", hass
)
refresh_runs = {
template_one: [],
template_five: [],
}
@ha.callback
def refresh_listener(event, updates):
for update in updates:
refresh_runs[update.template].append(update.result)
info = async_track_template_result(
hass,
[TrackTemplate(template_one, None), TrackTemplate(template_five, None)],
refresh_listener,
)
await hass.async_block_till_done()
info.async_refresh()
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0"]
assert refresh_runs[template_five] == ["0"]
hass.states.async_set("sensor.one", "any")
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0"]
assert refresh_runs[template_five] == ["0"]
info.async_refresh()
assert refresh_runs[template_one] == ["0", "1"]
assert refresh_runs[template_five] == ["0", "1"]
hass.states.async_set("sensor.two", "any")
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0", "1"]
assert refresh_runs[template_five] == ["0", "1"]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 1)
with patch(
"homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time
):
async_fire_time_changed(hass, next_time)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0", "1", "2"]
assert refresh_runs[template_five] == ["0", "1"]
hass.states.async_set("sensor.three", "any")
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0", "1", "2"]
assert refresh_runs[template_five] == ["0", "1"]
hass.states.async_set("sensor.four", "any")
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0", "1", "2"]
assert refresh_runs[template_five] == ["0", "1"]
hass.states.async_set("sensor.five", "any")
await hass.async_block_till_done()
assert refresh_runs[template_one] == ["0", "1", "2"]
assert refresh_runs[template_five] == ["0", "1"]
async def test_string(hass):
"""Test a string."""
template_refresh = Template("no_template", hass)
@ -1536,7 +1867,9 @@ async def test_async_track_template_result_multiple_templates_mixing_domain(hass
template_1 = Template("{{ states.switch.test.state == 'on' }}")
template_2 = Template("{{ states.switch.test.state == 'on' }}")
template_3 = Template("{{ states.switch.test.state == 'off' }}")
template_4 = Template("{{ states.switch | map(attribute='entity_id') | list }}")
template_4 = Template(
"{{ rate_limit(seconds=0) }}{{ states.switch | map(attribute='entity_id') | list }}"
)
refresh_runs = []

View File

@ -0,0 +1,108 @@
"""Tests for ratelimit."""
import asyncio
from datetime import timedelta
from homeassistant.core import callback
from homeassistant.helpers import ratelimit
from homeassistant.util import dt as dt_util
async def test_hit(hass):
"""Test hitting the rate limit."""
refresh_called = False
@callback
def _refresh():
nonlocal refresh_called
refresh_called = True
return
rate_limiter = ratelimit.KeyedRateLimit(hass)
rate_limiter.async_triggered("key1", dt_util.utcnow())
assert (
rate_limiter.async_schedule_action(
"key1", timedelta(seconds=0.001), dt_util.utcnow(), _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", timedelta(seconds=0.001), dt_util.utcnow(), _refresh
)
is None
)
rate_limiter.async_remove()
async def test_miss(hass):
"""Test missing the rate limit."""
refresh_called = False
@callback
def _refresh():
nonlocal refresh_called
refresh_called = True
return
rate_limiter = ratelimit.KeyedRateLimit(hass)
assert (
rate_limiter.async_schedule_action(
"key1", timedelta(seconds=0.1), dt_util.utcnow(), _refresh
)
is None
)
assert not refresh_called
assert not rate_limiter.async_has_timer("key1")
assert (
rate_limiter.async_schedule_action(
"key1", timedelta(seconds=0.1), dt_util.utcnow(), _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):
"""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
return
rate_limiter = ratelimit.KeyedRateLimit(hass)
rate_limiter.async_triggered("key1", dt_util.utcnow())
assert (
rate_limiter.async_schedule_action("key1", None, dt_util.utcnow(), _refresh)
is None
)
assert not refresh_called
assert not rate_limiter.async_has_timer("key1")
rate_limiter.async_triggered("key1", dt_util.utcnow())
assert (
rate_limiter.async_schedule_action("key1", None, dt_util.utcnow(), _refresh)
is None
)
assert not refresh_called
assert not rate_limiter.async_has_timer("key1")
rate_limiter.async_remove()

View File

@ -1,5 +1,5 @@
"""Test Home Assistant template helper methods."""
from datetime import datetime
from datetime import datetime, timedelta
import math
import random
@ -149,6 +149,7 @@ def test_iterating_all_states(hass):
info = render_to_info(hass, tmpl_str)
assert_result_info(info, "", all_states=True)
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
hass.states.async_set("test.object", "happy")
hass.states.async_set("sensor.temperature", 10)
@ -165,6 +166,7 @@ def test_iterating_domain_states(hass):
info = render_to_info(hass, tmpl_str)
assert_result_info(info, "", domains=["sensor"])
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
hass.states.async_set("test.object", "happy")
hass.states.async_set("sensor.back_door", "open")
@ -1342,6 +1344,7 @@ async def test_closest_function_home_vs_group_entity_id(hass):
assert_result_info(
info, "test_domain.object", {"group.location_group", "test_domain.object"}
)
assert info.rate_limit is None
async def test_closest_function_home_vs_group_state(hass):
@ -1369,20 +1372,24 @@ async def test_closest_function_home_vs_group_state(hass):
assert_result_info(
info, "test_domain.object", {"group.location_group", "test_domain.object"}
)
assert info.rate_limit is None
info = render_to_info(hass, "{{ closest(states.group.location_group).entity_id }}")
assert_result_info(
info, "test_domain.object", {"test_domain.object", "group.location_group"}
)
assert info.rate_limit is None
async def test_expand(hass):
"""Test expand function."""
info = render_to_info(hass, "{{ expand('test.object') }}")
assert_result_info(info, "[]", ["test.object"])
assert info.rate_limit is None
info = render_to_info(hass, "{{ expand(56) }}")
assert_result_info(info, "[]")
assert info.rate_limit is None
hass.states.async_set("test.object", "happy")
@ -1390,17 +1397,20 @@ async def test_expand(hass):
hass, "{{ expand('test.object') | map(attribute='entity_id') | join(', ') }}"
)
assert_result_info(info, "test.object", ["test.object"])
assert info.rate_limit is None
info = render_to_info(
hass,
"{{ expand('group.new_group') | map(attribute='entity_id') | join(', ') }}",
)
assert_result_info(info, "", ["group.new_group"])
assert info.rate_limit is None
info = render_to_info(
hass, "{{ expand(states.group) | map(attribute='entity_id') | join(', ') }}"
)
assert_result_info(info, "", [], ["group"])
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
assert await async_setup_component(hass, "group", {})
await hass.async_block_till_done()
@ -1411,6 +1421,7 @@ async def test_expand(hass):
"{{ expand('group.new_group') | map(attribute='entity_id') | join(', ') }}",
)
assert_result_info(info, "test.object", {"group.new_group", "test.object"})
assert info.rate_limit is None
info = render_to_info(
hass, "{{ expand(states.group) | map(attribute='entity_id') | join(', ') }}"
@ -1418,6 +1429,7 @@ async def test_expand(hass):
assert_result_info(
info, "test.object", {"test.object", "group.new_group"}, ["group"]
)
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
info = render_to_info(
hass,
@ -1432,6 +1444,7 @@ async def test_expand(hass):
" | map(attribute='entity_id') | join(', ') }}",
)
assert_result_info(info, "test.object", {"test.object", "group.new_group"})
assert info.rate_limit is None
hass.states.async_set("sensor.power_1", 0)
hass.states.async_set("sensor.power_2", 200.2)
@ -1452,6 +1465,7 @@ async def test_expand(hass):
str(200.2 + 400.4),
{"group.power_sensors", "sensor.power_1", "sensor.power_2", "sensor.power_3"},
)
assert info.rate_limit is None
def test_closest_function_to_coord(hass):
@ -1517,6 +1531,7 @@ def test_async_render_to_info_with_branching(hass):
""",
)
assert_result_info(info, "off", {"light.a", "light.c"})
assert info.rate_limit is None
info = render_to_info(
hass,
@ -1528,6 +1543,7 @@ def test_async_render_to_info_with_branching(hass):
""",
)
assert_result_info(info, "on", {"light.a", "light.b"})
assert info.rate_limit is None
def test_async_render_to_info_with_complex_branching(hass):
@ -1564,6 +1580,7 @@ def test_async_render_to_info_with_complex_branching(hass):
)
assert_result_info(info, "['sensor.a']", {"light.a", "light.b"}, {"sensor"})
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
async def test_async_render_to_info_with_wildcard_matching_entity_id(hass):
@ -1589,6 +1606,7 @@ async def test_async_render_to_info_with_wildcard_matching_entity_id(hass):
"cover.office_skylight",
}
assert info.all_states is True
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
async def test_async_render_to_info_with_wildcard_matching_state(hass):
@ -1619,6 +1637,7 @@ async def test_async_render_to_info_with_wildcard_matching_state(hass):
"cover.office_skylight",
}
assert info.all_states is True
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
hass.states.async_set("binary_sensor.door", "closed")
info = render_to_info(hass, template_complex_str)
@ -1632,6 +1651,7 @@ async def test_async_render_to_info_with_wildcard_matching_state(hass):
"cover.office_skylight",
}
assert info.all_states is True
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
template_cover_str = """
@ -1653,6 +1673,7 @@ async def test_async_render_to_info_with_wildcard_matching_state(hass):
"cover.office_skylight",
}
assert info.all_states is False
assert info.rate_limit == template.DEFAULT_RATE_LIMIT
def test_nested_async_render_to_info_case(hass):
@ -1665,6 +1686,7 @@ def test_nested_async_render_to_info_case(hass):
hass, "{{ states[states['input_select.picker'].state].state }}", {}
)
assert_result_info(info, "off", {"input_select.picker", "vacuum.a"})
assert info.rate_limit is None
def test_result_as_boolean(hass):
@ -2459,6 +2481,7 @@ async def test_lifecycle(hass):
info = tmp.async_render_to_info()
assert info.all_states is False
assert info.all_states_lifecycle is True
assert info.rate_limit is None
assert info.entities == set()
assert info.domains == set()
assert info.domains_lifecycle == set()
@ -2594,3 +2617,28 @@ async def test_unavailable_states(hass):
hass,
)
assert tpl.async_render() == "light.none, light.unavailable, light.unknown"
async def test_rate_limit(hass):
"""Test we can pickup a rate limit directive."""
tmp = template.Template("{{ states | count }}", hass)
info = tmp.async_render_to_info()
assert info.rate_limit is None
tmp = template.Template("{{ rate_limit(minutes=1) }}{{ states | count }}", hass)
info = tmp.async_render_to_info()
assert info.rate_limit == timedelta(minutes=1)
tmp = template.Template("{{ rate_limit(minutes=1) }}random", hass)
info = tmp.async_render_to_info()
assert info.result() == "random"
assert info.rate_limit == timedelta(minutes=1)
tmp = template.Template("{{ rate_limit(seconds=0) }}random", hass)
info = tmp.async_render_to_info()
assert info.result() == "random"
assert info.rate_limit == timedelta(seconds=0)