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