From d7cf05e693f5f2d6588a15460ba99350c67c59b2 Mon Sep 17 00:00:00 2001 From: Andy Castille Date: Wed, 11 Sep 2024 11:11:06 -0700 Subject: [PATCH] Allow attaching additional data to schedule helper blocks (#116585) * Add a new optional "data" key when defining time ranges for the schedule component that exposes the provided data in the state attributes of the schedule entity when that time range is active * Exclude all schedule entry custom data attributes from the recorder (with tests) * Fix setting schedule attributes to exclude from recorder, update test to verify the attributes exist but are not recorded * Fix test to ensure schedule data attributes are not recorded * Use vol.Any in place of vol.Or Co-authored-by: Erik Montnemery * Remove schedule block custom data shorthand as requested in https://github.com/home-assistant/core/pull/116585#pullrequestreview-2280260436 * Update homeassistant/components/schedule/__init__.py --------- Co-authored-by: Erik Montnemery --- homeassistant/components/schedule/__init__.py | 44 +++++++++++++++++-- homeassistant/components/schedule/const.py | 1 + tests/components/schedule/test_init.py | 33 +++++++++++--- tests/components/schedule/test_recorder.py | 35 +++++++++++++-- 4 files changed, 100 insertions(+), 13 deletions(-) diff --git a/homeassistant/components/schedule/__init__.py b/homeassistant/components/schedule/__init__.py index 08d0b083f7c..24ce4f3b3fa 100644 --- a/homeassistant/components/schedule/__init__.py +++ b/homeassistant/components/schedule/__init__.py @@ -39,6 +39,7 @@ from homeassistant.util import dt as dt_util from .const import ( ATTR_NEXT_EVENT, CONF_ALL_DAYS, + CONF_DATA, CONF_FROM, CONF_TO, DOMAIN, @@ -55,7 +56,7 @@ def valid_schedule(schedule: list[dict[str, str]]) -> list[dict[str, str]]: Ensure they have no overlap and the end time is greater than the start time. """ - # Emtpty schedule is valid + # Empty schedule is valid if not schedule: return schedule @@ -109,9 +110,13 @@ BASE_SCHEMA: VolDictType = { vol.Optional(CONF_ICON): cv.icon, } +# Extra data that the user can set on each time range +CUSTOM_DATA_SCHEMA = vol.Schema({str: vol.Any(bool, str, int, float)}) + TIME_RANGE_SCHEMA: VolDictType = { vol.Required(CONF_FROM): cv.time, vol.Required(CONF_TO): deserialize_to_time, + vol.Optional(CONF_DATA): CUSTOM_DATA_SCHEMA, } # Serialize time in validated config @@ -119,6 +124,7 @@ STORAGE_TIME_RANGE_SCHEMA = vol.Schema( { vol.Required(CONF_FROM): vol.Coerce(str), vol.Required(CONF_TO): serialize_to_time, + vol.Optional(CONF_DATA): CUSTOM_DATA_SCHEMA, } ) @@ -135,7 +141,6 @@ STORAGE_SCHEDULE_SCHEMA: VolDictType = { for day in CONF_ALL_DAYS } - # Validate YAML config CONFIG_SCHEMA = vol.Schema( {DOMAIN: cv.schema_with_slug_keys(vol.All(BASE_SCHEMA | SCHEDULE_SCHEMA))}, @@ -152,7 +157,7 @@ ENTITY_SCHEMA = vol.Schema( async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: - """Set up an input select.""" + """Set up a schedule.""" component = EntityComponent[Schedule](LOGGER, DOMAIN, hass) id_manager = IDManager() @@ -253,6 +258,12 @@ class Schedule(CollectionEntity): self._attr_name = self._config[CONF_NAME] self._attr_unique_id = self._config[CONF_ID] + # Exclude any custom attributes that may be present on time ranges from recording. + self._unrecorded_attributes = self.all_custom_data_keys() + self._Entity__combined_unrecorded_attributes = ( + self._entity_component_unrecorded_attributes | self._unrecorded_attributes + ) + @classmethod def from_storage(cls, config: ConfigType) -> Schedule: """Return entity instance initialized from storage.""" @@ -300,9 +311,11 @@ class Schedule(CollectionEntity): # Note that any time in the day is treated as smaller than time.max. if now.time() < time_range[CONF_TO] or time_range[CONF_TO] == time.max: self._attr_state = STATE_ON + current_data = time_range.get(CONF_DATA) break else: self._attr_state = STATE_OFF + current_data = None # Find next event in the schedule, loop over each day (starting with # the current day) until the next event has been found. @@ -344,6 +357,11 @@ class Schedule(CollectionEntity): self._attr_extra_state_attributes = { ATTR_NEXT_EVENT: next_event, } + + if current_data: + # Add each key/value pair in the data to the entity's state attributes + self._attr_extra_state_attributes.update(current_data) + self.async_write_ha_state() if next_event: @@ -352,3 +370,23 @@ class Schedule(CollectionEntity): self._update, next_event, ) + + def all_custom_data_keys(self) -> frozenset[str]: + """Return the set of all currently used custom data attribute keys.""" + data_keys = set() + + for weekday in WEEKDAY_TO_CONF.values(): + if not (weekday_config := self._config.get(weekday)): + continue # this weekday is not configured + + for time_range in weekday_config: + time_range_custom_data = time_range.get(CONF_DATA) + + if not time_range_custom_data or not isinstance( + time_range_custom_data, dict + ): + continue # this time range has no custom data, or it is not a dict + + data_keys.update(time_range_custom_data.keys()) + + return frozenset(data_keys) diff --git a/homeassistant/components/schedule/const.py b/homeassistant/components/schedule/const.py index 5ec57aae78d..6687dafefdb 100644 --- a/homeassistant/components/schedule/const.py +++ b/homeassistant/components/schedule/const.py @@ -6,6 +6,7 @@ from typing import Final DOMAIN: Final = "schedule" LOGGER = logging.getLogger(__package__) +CONF_DATA: Final = "data" CONF_FRIDAY: Final = "friday" CONF_FROM: Final = "from" CONF_MONDAY: Final = "monday" diff --git a/tests/components/schedule/test_init.py b/tests/components/schedule/test_init.py index 7cd59f19033..18346122bfd 100644 --- a/tests/components/schedule/test_init.py +++ b/tests/components/schedule/test_init.py @@ -12,6 +12,7 @@ import pytest from homeassistant.components.schedule import STORAGE_VERSION, STORAGE_VERSION_MINOR from homeassistant.components.schedule.const import ( ATTR_NEXT_EVENT, + CONF_DATA, CONF_FRIDAY, CONF_FROM, CONF_MONDAY, @@ -66,13 +67,21 @@ def schedule_setup( CONF_NAME: "from storage", CONF_ICON: "mdi:party-popper", CONF_FRIDAY: [ - {CONF_FROM: "17:00:00", CONF_TO: "23:59:59"}, + { + CONF_FROM: "17:00:00", + CONF_TO: "23:59:59", + CONF_DATA: {"party_level": "epic"}, + }, ], CONF_SATURDAY: [ {CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}, ], CONF_SUNDAY: [ - {CONF_FROM: "00:00:00", CONF_TO: "24:00:00"}, + { + CONF_FROM: "00:00:00", + CONF_TO: "24:00:00", + CONF_DATA: {"entry": "VIPs only"}, + }, ], } ] @@ -95,9 +104,21 @@ def schedule_setup( CONF_TUESDAY: [{CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}], CONF_WEDNESDAY: [{CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}], CONF_THURSDAY: [{CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}], - CONF_FRIDAY: [{CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}], + CONF_FRIDAY: [ + { + CONF_FROM: "00:00:00", + CONF_TO: "23:59:59", + CONF_DATA: {"party_level": "epic"}, + } + ], CONF_SATURDAY: [{CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}], - CONF_SUNDAY: [{CONF_FROM: "00:00:00", CONF_TO: "23:59:59"}], + CONF_SUNDAY: [ + { + CONF_FROM: "00:00:00", + CONF_TO: "23:59:59", + CONF_DATA: {"entry": "VIPs only"}, + } + ], } } } @@ -557,13 +578,13 @@ async def test_ws_list( assert len(result) == 1 assert result["from_storage"][ATTR_NAME] == "from storage" assert result["from_storage"][CONF_FRIDAY] == [ - {CONF_FROM: "17:00:00", CONF_TO: "23:59:59"} + {CONF_FROM: "17:00:00", CONF_TO: "23:59:59", CONF_DATA: {"party_level": "epic"}} ] assert result["from_storage"][CONF_SATURDAY] == [ {CONF_FROM: "00:00:00", CONF_TO: "23:59:59"} ] assert result["from_storage"][CONF_SUNDAY] == [ - {CONF_FROM: "00:00:00", CONF_TO: "24:00:00"} + {CONF_FROM: "00:00:00", CONF_TO: "24:00:00", CONF_DATA: {"entry": "VIPs only"}} ] assert "from_yaml" not in result diff --git a/tests/components/schedule/test_recorder.py b/tests/components/schedule/test_recorder.py index a7410472a44..85aef3e1990 100644 --- a/tests/components/schedule/test_recorder.py +++ b/tests/components/schedule/test_recorder.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import timedelta +from freezegun.api import FrozenDateTimeFactory import pytest from homeassistant.components.recorder.history import get_significant_states @@ -18,8 +19,11 @@ from tests.components.recorder.common import async_wait_recording_done @pytest.mark.usefixtures("recorder_mock", "enable_custom_integrations") -async def test_exclude_attributes(hass: HomeAssistant) -> None: +async def test_exclude_attributes( + hass: HomeAssistant, freezer: FrozenDateTimeFactory +) -> None: """Test attributes to be excluded.""" + freezer.move_to("2024-08-02 06:30:00-07:00") # Before Friday event now = dt_util.utcnow() assert await async_setup_component( hass, @@ -33,9 +37,13 @@ async def test_exclude_attributes(hass: HomeAssistant) -> None: "tuesday": [{"from": "2:00", "to": "3:00"}], "wednesday": [{"from": "3:00", "to": "4:00"}], "thursday": [{"from": "5:00", "to": "6:00"}], - "friday": [{"from": "7:00", "to": "8:00"}], + "friday": [ + {"from": "7:00", "to": "8:00", "data": {"party_level": "epic"}} + ], "saturday": [{"from": "9:00", "to": "10:00"}], - "sunday": [{"from": "11:00", "to": "12:00"}], + "sunday": [ + {"from": "11:00", "to": "12:00", "data": {"entry": "VIPs only"}} + ], } } }, @@ -48,8 +56,25 @@ async def test_exclude_attributes(hass: HomeAssistant) -> None: assert state.attributes[ATTR_ICON] assert state.attributes[ATTR_NEXT_EVENT] + # Move to during Friday event + freezer.move_to("2024-08-02 07:30:00-07:00") + async_fire_time_changed(hass, fire_all=True) await hass.async_block_till_done() - async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=5)) + state = hass.states.get("schedule.test") + assert "entry" not in state.attributes + assert state.attributes["party_level"] == "epic" + + # Move to during Sunday event + freezer.move_to("2024-08-04 11:30:00-07:00") + async_fire_time_changed(hass, fire_all=True) + await hass.async_block_till_done() + state = hass.states.get("schedule.test") + assert "party_level" not in state.attributes + assert state.attributes["entry"] == "VIPs only" + + await hass.async_block_till_done() + freezer.tick(timedelta(minutes=5)) + async_fire_time_changed(hass) await hass.async_block_till_done() await async_wait_recording_done(hass) @@ -63,3 +88,5 @@ async def test_exclude_attributes(hass: HomeAssistant) -> None: assert ATTR_FRIENDLY_NAME in state.attributes assert ATTR_ICON in state.attributes assert ATTR_NEXT_EVENT not in state.attributes + assert "entry" not in state.attributes + assert "party_level" not in state.attributes