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 <erik@montnemery.com>

* 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 <erik@montnemery.com>
This commit is contained in:
Andy Castille 2024-09-11 11:11:06 -07:00 committed by GitHub
parent 75d3ea34fc
commit d7cf05e693
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 100 additions and 13 deletions

View File

@ -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)

View File

@ -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"

View File

@ -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

View File

@ -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