From 3a9b2392f895f636994c987b92a3d7da6f522f70 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 19 Oct 2020 03:17:51 -0500
Subject: [PATCH] Suppress domain and all listeners during template rate limit
 (#42005)

---
 homeassistant/helpers/event.py | 24 +++++++++++--
 tests/helpers/test_event.py    | 65 ++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index 044f18c4255..a3eb981140d 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -1,5 +1,6 @@
 """Helpers for listening to events."""
 import asyncio
+import copy
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 import functools as ft
@@ -820,6 +821,8 @@ class _TrackTemplateResultInfo:
             if not _event_triggers_rerender(event, info):
                 return False
 
+            had_timer = self._rate_limit.async_has_timer(template)
+
             if self._rate_limit.async_schedule_action(
                 template,
                 _rate_limit_for_event(event, info, track_template_),
@@ -829,7 +832,7 @@ class _TrackTemplateResultInfo:
                 (track_template_,),
                 True,
             ):
-                return False
+                return not had_timer
 
             _LOGGER.debug(
                 "Template update %s triggered by event: %s",
@@ -893,7 +896,14 @@ class _TrackTemplateResultInfo:
         if info_changed:
             assert self._track_state_changes
             self._track_state_changes.async_update_listeners(
-                _render_infos_to_track_states(self._info.values()),
+                _render_infos_to_track_states(
+                    [
+                        _suppress_domain_all_in_render_info(self._info[template])
+                        if self._rate_limit.async_has_timer(template)
+                        else self._info[template]
+                        for template in self._info
+                    ]
+                )
             )
             _LOGGER.debug(
                 "Template group %s listens for %s",
@@ -1458,3 +1468,13 @@ def _rate_limit_for_event(
 
     rate_limit: Optional[timedelta] = info.rate_limit
     return rate_limit
+
+
+def _suppress_domain_all_in_render_info(render_info: RenderInfo) -> RenderInfo:
+    """Remove the domains and all_states from render info during a ratelimit."""
+    rate_limited_render_info = copy.copy(render_info)
+    rate_limited_render_info.all_states = False
+    rate_limited_render_info.all_states_lifecycle = False
+    rate_limited_render_info.domains = set()
+    rate_limited_render_info.domains_lifecycle = set()
+    return rate_limited_render_info
diff --git a/tests/helpers/test_event.py b/tests/helpers/test_event.py
index f63abfa2ed2..0a4b85edcf4 100644
--- a/tests/helpers/test_event.py
+++ b/tests/helpers/test_event.py
@@ -1494,6 +1494,71 @@ async def test_track_template_rate_limit(hass):
     assert refresh_runs == [0, 1, 2, 4]
 
 
+async def test_track_template_rate_limit_suppress_listener(hass):
+    """Test template rate limit will suppress the listener during the 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()
+
+    assert info.listeners == {"all": True, "domains": set(), "entities": set()}
+    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()
+    # Should be suppressed during the rate limit
+    assert info.listeners == {"all": False, "domains": set(), "entities": set()}
+    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()
+    # Rate limit released and the all listener returns
+    assert info.listeners == {"all": True, "domains": set(), "entities": set()}
+    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]
+    # Rate limit hit and the all listener is shut off
+    assert info.listeners == {"all": False, "domains": set(), "entities": set()}
+    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()
+    # Rate limit released and the all listener returns
+    assert info.listeners == {"all": True, "domains": set(), "entities": set()}
+    assert refresh_runs == [0, 1, 2, 4]
+    hass.states.async_set("sensor.five", "any")
+    await hass.async_block_till_done()
+    # Rate limit hit and the all listener is shut off
+    assert info.listeners == {"all": False, "domains": set(), "entities": set()}
+    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)