From fa9866db96ea2e63cf05d2c3c97e8bdc64b5f604 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Thu, 30 Jul 2020 02:51:30 +0200
Subject: [PATCH] Add support for multiple time triggers in automations
 (#37975)

* Add support for multiple time triggers in automations

* Attach with single callback

* Patch time in tests

* Improve test coverage

* Adjusting my facepalm moment
---
 homeassistant/components/automation/time.py | 29 ++++++--
 tests/components/automation/test_time.py    | 82 ++++++++++++++++++++-
 2 files changed, 104 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/automation/time.py b/homeassistant/components/automation/time.py
index 5f461952960..f59ceff81ea 100644
--- a/homeassistant/components/automation/time.py
+++ b/homeassistant/components/automation/time.py
@@ -13,20 +13,37 @@ from homeassistant.helpers.event import async_track_time_change
 _LOGGER = logging.getLogger(__name__)
 
 TRIGGER_SCHEMA = vol.Schema(
-    {vol.Required(CONF_PLATFORM): "time", vol.Required(CONF_AT): cv.time}
+    {
+        vol.Required(CONF_PLATFORM): "time",
+        vol.Required(CONF_AT): vol.All(cv.ensure_list, [cv.time]),
+    }
 )
 
 
 async def async_attach_trigger(hass, config, action, automation_info):
     """Listen for state changes based on configuration."""
-    at_time = config.get(CONF_AT)
-    hours, minutes, seconds = at_time.hour, at_time.minute, at_time.second
+    at_times = config[CONF_AT]
 
     @callback
     def time_automation_listener(now):
         """Listen for time changes and calls action."""
         hass.async_run_job(action, {"trigger": {"platform": "time", "now": now}})
 
-    return async_track_time_change(
-        hass, time_automation_listener, hour=hours, minute=minutes, second=seconds
-    )
+    removes = [
+        async_track_time_change(
+            hass,
+            time_automation_listener,
+            hour=at_time.hour,
+            minute=at_time.minute,
+            second=at_time.second,
+        )
+        for at_time in at_times
+    ]
+
+    @callback
+    def remove_track_time_changes():
+        """Remove tracked time changes."""
+        for remove in removes:
+            remove()
+
+    return remove_track_time_changes
diff --git a/tests/components/automation/test_time.py b/tests/components/automation/test_time.py
index c93cdbc36e9..c8b95985636 100644
--- a/tests/components/automation/test_time.py
+++ b/tests/components/automation/test_time.py
@@ -4,10 +4,11 @@ from datetime import timedelta
 import pytest
 
 import homeassistant.components.automation as automation
+from homeassistant.const import ATTR_ENTITY_ID, SERVICE_TURN_OFF
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from tests.async_mock import patch
+from tests.async_mock import Mock, patch
 from tests.common import (
     assert_setup_component,
     async_fire_time_changed,
@@ -66,6 +67,53 @@ async def test_if_fires_using_at(hass, calls):
     assert calls[0].data["some"] == "time - 5"
 
 
+async def test_if_fires_using_multiple_at(hass, calls):
+    """Test for firing at."""
+
+    now = dt_util.utcnow()
+
+    time_that_will_not_match_right_away = now.replace(
+        year=now.year + 1, hour=4, minute=59, second=0
+    )
+
+    with patch(
+        "homeassistant.util.dt.utcnow", return_value=time_that_will_not_match_right_away
+    ):
+        assert await async_setup_component(
+            hass,
+            automation.DOMAIN,
+            {
+                automation.DOMAIN: {
+                    "trigger": {"platform": "time", "at": ["5:00:00", "6:00:00"]},
+                    "action": {
+                        "service": "test.automation",
+                        "data_template": {
+                            "some": "{{ trigger.platform }} - {{ trigger.now.hour }}"
+                        },
+                    },
+                }
+            },
+        )
+
+    now = dt_util.utcnow()
+
+    async_fire_time_changed(
+        hass, now.replace(year=now.year + 1, hour=5, minute=0, second=0)
+    )
+
+    await hass.async_block_till_done()
+    assert len(calls) == 1
+    assert calls[0].data["some"] == "time - 5"
+
+    async_fire_time_changed(
+        hass, now.replace(year=now.year + 1, hour=6, minute=0, second=0)
+    )
+
+    await hass.async_block_till_done()
+    assert len(calls) == 2
+    assert calls[1].data["some"] == "time - 6"
+
+
 async def test_if_not_fires_using_wrong_at(hass, calls):
     """YAML translates time values to total seconds.
 
@@ -231,3 +279,35 @@ async def test_if_action_list_weekday(hass, calls):
         await hass.async_block_till_done()
 
     assert len(calls) == 2
+
+
+async def test_untrack_time_change(hass):
+    """Test for removing tracked time changes."""
+    mock_track_time_change = Mock()
+    with patch(
+        "homeassistant.components.automation.time.async_track_time_change",
+        return_value=mock_track_time_change,
+    ):
+        assert await async_setup_component(
+            hass,
+            automation.DOMAIN,
+            {
+                automation.DOMAIN: {
+                    "alias": "test",
+                    "trigger": {
+                        "platform": "time",
+                        "at": ["5:00:00", "6:00:00", "7:00:00"],
+                    },
+                    "action": {"service": "test.automation", "data": {"test": "test"}},
+                }
+            },
+        )
+
+    await hass.services.async_call(
+        automation.DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: "automation.test"},
+        blocking=True,
+    )
+
+    assert len(mock_track_time_change.mock_calls) == 3