From 51da605b9f35d29e9b00fbf9a4eae0aca9123c3c Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 4 Oct 2020 15:40:04 -0500
Subject: [PATCH] Remove manual rate_limit control directive from templates
 (#41225)

Increase default rate limit for all states and entire
domain states to one minute

Ensure specifically referenced entities are excluded from
the rate limit
---
 homeassistant/helpers/event.py    |  21 +++-
 homeassistant/helpers/template.py |  25 +---
 tests/helpers/test_event.py       | 196 +++++-------------------------
 tests/helpers/test_template.py    |  27 +---
 4 files changed, 51 insertions(+), 218 deletions(-)

diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index d143ed1d106..b6d59bb500c 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -802,7 +802,7 @@ class _TrackTemplateResultInfo:
 
             if self._rate_limit.async_schedule_action(
                 template,
-                info.rate_limit or track_template_.rate_limit,
+                _rate_limit_for_event(event, info, track_template_),
                 now,
                 self._refresh,
                 event,
@@ -1376,3 +1376,22 @@ def _event_triggers_rerender(event: Event, info: RenderInfo) -> bool:
         return False
 
     return bool(info.filter_lifecycle(entity_id))
+
+
+@callback
+def _rate_limit_for_event(
+    event: Event, info: RenderInfo, track_template_: TrackTemplate
+) -> Optional[timedelta]:
+    """Determine the rate limit for an event."""
+    entity_id = event.data.get(ATTR_ENTITY_ID)
+
+    # Specifically referenced entities are excluded
+    # from the rate limit
+    if entity_id in info.entities:
+        return None
+
+    if track_template_.rate_limit is not None:
+        return track_template_.rate_limit
+
+    rate_limit: Optional[timedelta] = info.rate_limit
+    return rate_limit
diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py
index b877e0b0e12..c0ceb17fdd7 100644
--- a/homeassistant/helpers/template.py
+++ b/homeassistant/helpers/template.py
@@ -72,7 +72,7 @@ _COLLECTABLE_STATE_ATTRIBUTES = {
     "name",
 }
 
-DEFAULT_RATE_LIMIT = timedelta(seconds=1)
+DEFAULT_RATE_LIMIT = timedelta(minutes=1)
 
 
 @bind_hass
@@ -489,26 +489,6 @@ 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."""
 
@@ -1310,11 +1290,10 @@ 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, RateLimit)) or super().is_safe_callable(obj)
+        return isinstance(obj, AllStates) or super().is_safe_callable(obj)
 
     def is_safe_attribute(self, obj, attr, value):
         """Test if attribute is safe."""
diff --git a/tests/helpers/test_event.py b/tests/helpers/test_event.py
index 32d203ddc6e..bf0efaf8cf6 100644
--- a/tests/helpers/test_event.py
+++ b/tests/helpers/test_event.py
@@ -927,7 +927,6 @@ 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" %}
@@ -948,7 +947,9 @@ async def test_track_template_result_complex(hass):
     hass.states.async_set("lock.one", "locked")
 
     info = async_track_template_result(
-        hass, [TrackTemplate(template_complex, None)], specific_run_callback
+        hass,
+        [TrackTemplate(template_complex, None, timedelta(seconds=0))],
+        specific_run_callback,
     )
     await hass.async_block_till_done()
 
@@ -1236,7 +1237,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 }},
@@ -1246,6 +1247,7 @@ async def test_track_template_result_iterator(hass):
                     hass,
                 ),
                 None,
+                timedelta(seconds=0),
             )
         ],
         iterator_callback,
@@ -1268,11 +1270,12 @@ async def test_track_template_result_iterator(hass):
         [
             TrackTemplate(
                 Template(
-                    """{{ rate_limit(seconds=0) }}{{ states.sensor|selectattr("state","equalto","on")
+                    """{{ states.sensor|selectattr("state","equalto","on")
                 |join(",", attribute="entity_id") }}""",
                     hass,
                 ),
                 None,
+                timedelta(seconds=0),
             )
         ],
         filter_callback,
@@ -1452,62 +1455,6 @@ async def test_track_template_rate_limit(hass):
     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)
@@ -1541,19 +1488,11 @@ async def test_track_template_rate_limit_five(hass):
     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,
-    )
+async def test_specifically_referenced_entity_is_not_rate_limited(hass):
+    """Test template rate limit of 5 seconds."""
+    hass.states.async_set("sensor.one", "none")
+
+    template_refresh = Template('{{ states | count }}_{{ states("sensor.one") }}', hass)
 
     refresh_runs = []
 
@@ -1563,114 +1502,34 @@ async def test_track_template_rate_limit_changes(hass):
 
     info = async_track_template_result(
         hass,
-        [TrackTemplate(template_refresh, None)],
+        [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"]
+    assert refresh_runs == ["1_none"]
     hass.states.async_set("sensor.one", "any")
     await hass.async_block_till_done()
-    assert refresh_runs == ["0"]
+    assert refresh_runs == ["1_none", "1_any"]
     info.async_refresh()
-    assert refresh_runs == ["0", "1"]
+    assert refresh_runs == ["1_none", "1_any"]
     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"]
+    assert refresh_runs == ["1_none", "1_any"]
     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")
+    assert refresh_runs == ["1_none", "1_any"]
+    hass.states.async_set("sensor.one", "none")
     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"]
+    assert refresh_runs == ["1_none", "1_any", "3_none"]
 
 
 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
-    )
+    template_one = Template("{{ states | count }} ", hass)
+    template_five = Template("{{ states | count }}", hass)
 
     refresh_runs = {
         template_one: [],
@@ -1684,7 +1543,10 @@ async def test_track_two_templates_with_different_rate_limits(hass):
 
     info = async_track_template_result(
         hass,
-        [TrackTemplate(template_one, None), TrackTemplate(template_five, None)],
+        [
+            TrackTemplate(template_one, None, timedelta(seconds=0.1)),
+            TrackTemplate(template_five, None, timedelta(seconds=5)),
+        ],
         refresh_listener,
     )
 
@@ -1867,9 +1729,7 @@ 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(
-        "{{ rate_limit(seconds=0) }}{{ states.switch | map(attribute='entity_id') | list }}"
-    )
+    template_4 = Template("{{ states.switch | map(attribute='entity_id') | list }}")
 
     refresh_runs = []
 
@@ -1883,7 +1743,7 @@ async def test_async_track_template_result_multiple_templates_mixing_domain(hass
             TrackTemplate(template_1, None),
             TrackTemplate(template_2, None),
             TrackTemplate(template_3, None),
-            TrackTemplate(template_4, None),
+            TrackTemplate(template_4, None, timedelta(seconds=0)),
         ],
         refresh_listener,
     )
diff --git a/tests/helpers/test_template.py b/tests/helpers/test_template.py
index 0a5735a98ff..7d39675bdc0 100644
--- a/tests/helpers/test_template.py
+++ b/tests/helpers/test_template.py
@@ -1,5 +1,5 @@
 """Test Home Assistant template helper methods."""
-from datetime import datetime, timedelta
+from datetime import datetime
 import math
 import random
 
@@ -2617,28 +2617,3 @@ 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)