Use storage based collections for Timer platform (#30765)

* Use config dict for timer entity.

* Manage timer entities using collection helpers.

* Add tests.

* Make Timer duration JSON serializable.
This commit is contained in:
Alexei Chetroi 2020-01-14 18:32:48 -05:00 committed by Paulus Schoutsen
parent 6f84723fec
commit 0a7feba855
2 changed files with 309 additions and 50 deletions

View File

@ -1,15 +1,26 @@
"""Support for Timers.""" """Support for Timers."""
from datetime import timedelta from datetime import timedelta
import logging import logging
import typing
import voluptuous as vol import voluptuous as vol
from homeassistant.const import CONF_ICON, CONF_NAME, SERVICE_RELOAD from homeassistant.const import (
ATTR_EDITABLE,
CONF_ICON,
CONF_ID,
CONF_NAME,
SERVICE_RELOAD,
)
from homeassistant.core import callback
from homeassistant.helpers import collection, entity_registry
import homeassistant.helpers.config_validation as cv import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.entity_component import EntityComponent
from homeassistant.helpers.event import async_track_point_in_utc_time from homeassistant.helpers.event import async_track_point_in_utc_time
from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.restore_state import RestoreEntity
import homeassistant.helpers.service import homeassistant.helpers.service
from homeassistant.helpers.storage import Store
from homeassistant.helpers.typing import ConfigType, HomeAssistantType, ServiceCallType
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -17,7 +28,7 @@ _LOGGER = logging.getLogger(__name__)
DOMAIN = "timer" DOMAIN = "timer"
ENTITY_ID_FORMAT = DOMAIN + ".{}" ENTITY_ID_FORMAT = DOMAIN + ".{}"
DEFAULT_DURATION = timedelta(0) DEFAULT_DURATION = 0
ATTR_DURATION = "duration" ATTR_DURATION = "duration"
ATTR_REMAINING = "remaining" ATTR_REMAINING = "remaining"
CONF_DURATION = "duration" CONF_DURATION = "duration"
@ -37,6 +48,21 @@ SERVICE_PAUSE = "pause"
SERVICE_CANCEL = "cancel" SERVICE_CANCEL = "cancel"
SERVICE_FINISH = "finish" SERVICE_FINISH = "finish"
STORAGE_KEY = DOMAIN
STORAGE_VERSION = 1
CREATE_FIELDS = {
vol.Required(CONF_NAME): vol.All(str, vol.Length(min=1)),
vol.Optional(CONF_NAME): cv.string,
vol.Optional(CONF_ICON): cv.icon,
vol.Optional(CONF_DURATION, default=DEFAULT_DURATION): cv.time_period,
}
UPDATE_FIELDS = {
vol.Optional(CONF_NAME): cv.string,
vol.Optional(CONF_ICON): cv.icon,
vol.Optional(CONF_DURATION): cv.time_period,
}
def _none_to_empty_dict(value): def _none_to_empty_dict(value):
if value is None: if value is None:
@ -65,20 +91,55 @@ CONFIG_SCHEMA = vol.Schema(
RELOAD_SERVICE_SCHEMA = vol.Schema({}) RELOAD_SERVICE_SCHEMA = vol.Schema({})
async def async_setup(hass, config): async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
"""Set up a timer.""" """Set up an input select."""
component = EntityComponent(_LOGGER, DOMAIN, hass) component = EntityComponent(_LOGGER, DOMAIN, hass)
id_manager = collection.IDManager()
entities = await _async_process_config(hass, config) yaml_collection = collection.YamlCollection(
logging.getLogger(f"{__name__}.yaml_collection"), id_manager
)
collection.attach_entity_component_collection(
component, yaml_collection, Timer.from_yaml
)
async def reload_service_handler(service_call): storage_collection = TimerStorageCollection(
"""Remove all input booleans and load new ones from config.""" Store(hass, STORAGE_VERSION, STORAGE_KEY),
conf = await component.async_prepare_reload() logging.getLogger(f"{__name__}.storage_collection"),
if conf is None: id_manager,
)
collection.attach_entity_component_collection(component, storage_collection, Timer)
await yaml_collection.async_load(
[{CONF_ID: id_, **cfg} for id_, cfg in config[DOMAIN].items()]
)
await storage_collection.async_load()
collection.StorageCollectionWebsocket(
storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
).async_setup(hass)
async def _collection_changed(
change_type: str, item_id: str, config: typing.Optional[typing.Dict]
) -> None:
"""Handle a collection change: clean up entity registry on removals."""
if change_type != collection.CHANGE_REMOVED:
return return
new_entities = await _async_process_config(hass, conf)
if new_entities: ent_reg = await entity_registry.async_get_registry(hass)
await component.async_add_entities(new_entities) ent_reg.async_remove(ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id))
yaml_collection.async_add_listener(_collection_changed)
storage_collection.async_add_listener(_collection_changed)
async def reload_service_handler(service_call: ServiceCallType) -> None:
"""Reload yaml entities."""
conf = await component.async_prepare_reload(skip_reset=True)
if conf is None:
conf = {DOMAIN: {}}
await yaml_collection.async_load(
[{CONF_ID: id_, **cfg} for id_, cfg in conf[DOMAIN].items()]
)
homeassistant.helpers.service.async_register_admin_service( homeassistant.helpers.service.async_register_admin_service(
hass, hass,
@ -96,43 +157,55 @@ async def async_setup(hass, config):
component.async_register_entity_service(SERVICE_CANCEL, {}, "async_cancel") component.async_register_entity_service(SERVICE_CANCEL, {}, "async_cancel")
component.async_register_entity_service(SERVICE_FINISH, {}, "async_finish") component.async_register_entity_service(SERVICE_FINISH, {}, "async_finish")
if entities:
await component.async_add_entities(entities)
return True return True
async def _async_process_config(hass, config): class TimerStorageCollection(collection.StorageCollection):
"""Process config and create list of entities.""" """Timer storage based collection."""
entities = []
for object_id, cfg in config[DOMAIN].items(): CREATE_SCHEMA = vol.Schema(CREATE_FIELDS)
if not cfg: UPDATE_SCHEMA = vol.Schema(UPDATE_FIELDS)
cfg = {}
name = cfg.get(CONF_NAME) async def _process_create_data(self, data: typing.Dict) -> typing.Dict:
icon = cfg.get(CONF_ICON) """Validate the config is valid."""
duration = cfg[CONF_DURATION] data = self.CREATE_SCHEMA(data)
# make duration JSON serializeable
data[CONF_DURATION] = str(data[CONF_DURATION])
return data
entities.append(Timer(hass, object_id, name, icon, duration)) @callback
def _get_suggested_id(self, info: typing.Dict) -> str:
"""Suggest an ID based on the config."""
return info[CONF_NAME]
return entities async def _update_data(self, data: dict, update_data: typing.Dict) -> typing.Dict:
"""Return a new updated data object."""
data = {**data, **self.UPDATE_SCHEMA(update_data)}
# make duration JSON serializeable
data[CONF_DURATION] = str(data[CONF_DURATION])
return data
class Timer(RestoreEntity): class Timer(RestoreEntity):
"""Representation of a timer.""" """Representation of a timer."""
def __init__(self, hass, object_id, name, icon, duration): def __init__(self, config: typing.Dict):
"""Initialize a timer.""" """Initialize a timer."""
self.entity_id = ENTITY_ID_FORMAT.format(object_id) self._config = config
self._name = name self.editable = True
self._state = STATUS_IDLE self._state = STATUS_IDLE
self._duration = duration self._remaining = config[CONF_DURATION]
self._remaining = self._duration
self._icon = icon
self._hass = hass
self._end = None self._end = None
self._listener = None self._listener = None
@classmethod
def from_yaml(cls, config: typing.Dict) -> "Timer":
"""Return entity instance initialized from yaml storage."""
timer = cls(config)
timer.entity_id = ENTITY_ID_FORMAT.format(config[CONF_ID])
timer.editable = False
return timer
@property @property
def should_poll(self): def should_poll(self):
"""If entity should be polled.""" """If entity should be polled."""
@ -141,12 +214,12 @@ class Timer(RestoreEntity):
@property @property
def name(self): def name(self):
"""Return name of the timer.""" """Return name of the timer."""
return self._name return self._config.get(CONF_NAME)
@property @property
def icon(self): def icon(self):
"""Return the icon to be used for this entity.""" """Return the icon to be used for this entity."""
return self._icon return self._config.get(CONF_ICON)
@property @property
def state(self): def state(self):
@ -157,10 +230,16 @@ class Timer(RestoreEntity):
def state_attributes(self): def state_attributes(self):
"""Return the state attributes.""" """Return the state attributes."""
return { return {
ATTR_DURATION: str(self._duration), ATTR_DURATION: str(self._config[CONF_DURATION]),
ATTR_EDITABLE: self.editable,
ATTR_REMAINING: str(self._remaining), ATTR_REMAINING: str(self._remaining),
} }
@property
def unique_id(self) -> typing.Optional[str]:
"""Return unique id for the entity."""
return self._config[CONF_ID]
async def async_added_to_hass(self): async def async_added_to_hass(self):
"""Call when entity is about to be added to Home Assistant.""" """Call when entity is about to be added to Home Assistant."""
# If not None, we got an initial value. # If not None, we got an initial value.
@ -189,18 +268,18 @@ class Timer(RestoreEntity):
self._end = start + self._remaining self._end = start + self._remaining
else: else:
if newduration: if newduration:
self._duration = newduration self._config[CONF_DURATION] = newduration
self._remaining = newduration self._remaining = newduration
else: else:
self._remaining = self._duration self._remaining = self._config[CONF_DURATION]
self._end = start + self._duration self._end = start + self._config[CONF_DURATION]
self._hass.bus.async_fire(event, {"entity_id": self.entity_id}) self.hass.bus.async_fire(event, {"entity_id": self.entity_id})
self._listener = async_track_point_in_utc_time( self._listener = async_track_point_in_utc_time(
self._hass, self.async_finished, self._end self.hass, self.async_finished, self._end
) )
await self.async_update_ha_state() self.async_write_ha_state()
async def async_pause(self): async def async_pause(self):
"""Pause a timer.""" """Pause a timer."""
@ -212,8 +291,8 @@ class Timer(RestoreEntity):
self._remaining = self._end - dt_util.utcnow().replace(microsecond=0) self._remaining = self._end - dt_util.utcnow().replace(microsecond=0)
self._state = STATUS_PAUSED self._state = STATUS_PAUSED
self._end = None self._end = None
self._hass.bus.async_fire(EVENT_TIMER_PAUSED, {"entity_id": self.entity_id}) self.hass.bus.async_fire(EVENT_TIMER_PAUSED, {"entity_id": self.entity_id})
await self.async_update_ha_state() self.async_write_ha_state()
async def async_cancel(self): async def async_cancel(self):
"""Cancel a timer.""" """Cancel a timer."""
@ -223,8 +302,8 @@ class Timer(RestoreEntity):
self._state = STATUS_IDLE self._state = STATUS_IDLE
self._end = None self._end = None
self._remaining = timedelta() self._remaining = timedelta()
self._hass.bus.async_fire(EVENT_TIMER_CANCELLED, {"entity_id": self.entity_id}) self.hass.bus.async_fire(EVENT_TIMER_CANCELLED, {"entity_id": self.entity_id})
await self.async_update_ha_state() self.async_write_ha_state()
async def async_finish(self): async def async_finish(self):
"""Reset and updates the states, fire finished event.""" """Reset and updates the states, fire finished event."""
@ -234,8 +313,8 @@ class Timer(RestoreEntity):
self._listener = None self._listener = None
self._state = STATUS_IDLE self._state = STATUS_IDLE
self._remaining = timedelta() self._remaining = timedelta()
self._hass.bus.async_fire(EVENT_TIMER_FINISHED, {"entity_id": self.entity_id}) self.hass.bus.async_fire(EVENT_TIMER_FINISHED, {"entity_id": self.entity_id})
await self.async_update_ha_state() self.async_write_ha_state()
async def async_finished(self, time): async def async_finished(self, time):
"""Reset and updates the states, fire finished event.""" """Reset and updates the states, fire finished event."""
@ -245,5 +324,10 @@ class Timer(RestoreEntity):
self._listener = None self._listener = None
self._state = STATUS_IDLE self._state = STATUS_IDLE
self._remaining = timedelta() self._remaining = timedelta()
self._hass.bus.async_fire(EVENT_TIMER_FINISHED, {"entity_id": self.entity_id}) self.hass.bus.async_fire(EVENT_TIMER_FINISHED, {"entity_id": self.entity_id})
await self.async_update_ha_state() self.async_write_ha_state()
async def async_update_config(self, config: typing.Dict) -> None:
"""Handle when the config is updated."""
self._config = config
self.async_write_ha_state()

View File

@ -27,13 +27,17 @@ from homeassistant.components.timer import (
STATUS_PAUSED, STATUS_PAUSED,
) )
from homeassistant.const import ( from homeassistant.const import (
ATTR_EDITABLE,
ATTR_FRIENDLY_NAME, ATTR_FRIENDLY_NAME,
ATTR_ICON, ATTR_ICON,
ATTR_ID,
ATTR_NAME,
CONF_ENTITY_ID, CONF_ENTITY_ID,
SERVICE_RELOAD, SERVICE_RELOAD,
) )
from homeassistant.core import Context, CoreState from homeassistant.core import Context, CoreState
from homeassistant.exceptions import Unauthorized from homeassistant.exceptions import Unauthorized
from homeassistant.helpers import config_validation as cv, entity_registry
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
from homeassistant.util.dt import utcnow from homeassistant.util.dt import utcnow
@ -42,6 +46,38 @@ from tests.common import async_fire_time_changed
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@pytest.fixture
def storage_setup(hass, hass_storage):
"""Storage setup."""
async def _storage(items=None, config=None):
if items is None:
hass_storage[DOMAIN] = {
"key": DOMAIN,
"version": 1,
"data": {
"items": [
{
ATTR_ID: "from_storage",
ATTR_NAME: "timer from storage",
ATTR_DURATION: 0,
}
]
},
}
else:
hass_storage[DOMAIN] = {
"key": DOMAIN,
"version": 1,
"data": {"items": items},
}
if config is None:
config = {DOMAIN: {}}
return await async_setup_component(hass, DOMAIN, config)
return _storage
async def test_config(hass): async def test_config(hass):
"""Test config.""" """Test config."""
invalid_configs = [None, 1, {}, {"name with space": None}] invalid_configs = [None, 1, {}, {"name with space": None}]
@ -92,7 +128,9 @@ async def test_config_options(hass):
assert "0:00:10" == state_2.attributes.get(ATTR_DURATION) assert "0:00:10" == state_2.attributes.get(ATTR_DURATION)
assert STATUS_IDLE == state_3.state assert STATUS_IDLE == state_3.state
assert str(DEFAULT_DURATION) == state_3.attributes.get(CONF_DURATION) assert str(cv.time_period(DEFAULT_DURATION)) == state_3.attributes.get(
CONF_DURATION
)
async def test_methods_and_events(hass): async def test_methods_and_events(hass):
@ -208,6 +246,7 @@ async def test_no_initial_state_and_no_restore_state(hass):
async def test_config_reload(hass, hass_admin_user, hass_read_only_user): async def test_config_reload(hass, hass_admin_user, hass_read_only_user):
"""Test reload service.""" """Test reload service."""
count_start = len(hass.states.async_entity_ids()) count_start = len(hass.states.async_entity_ids())
ent_reg = await entity_registry.async_get_registry(hass)
_LOGGER.debug("ENTITIES @ start: %s", hass.states.async_entity_ids()) _LOGGER.debug("ENTITIES @ start: %s", hass.states.async_entity_ids())
@ -235,6 +274,9 @@ async def test_config_reload(hass, hass_admin_user, hass_read_only_user):
assert state_1 is not None assert state_1 is not None
assert state_2 is not None assert state_2 is not None
assert state_3 is None assert state_3 is None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "test_1") is not None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "test_2") is not None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "test_3") is None
assert STATUS_IDLE == state_1.state assert STATUS_IDLE == state_1.state
assert ATTR_ICON not in state_1.attributes assert ATTR_ICON not in state_1.attributes
@ -283,6 +325,9 @@ async def test_config_reload(hass, hass_admin_user, hass_read_only_user):
assert state_1 is None assert state_1 is None
assert state_2 is not None assert state_2 is not None
assert state_3 is not None assert state_3 is not None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "test_1") is None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "test_2") is not None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, "test_3") is not None
assert STATUS_IDLE == state_2.state assert STATUS_IDLE == state_2.state
assert "Hello World reloaded" == state_2.attributes.get(ATTR_FRIENDLY_NAME) assert "Hello World reloaded" == state_2.attributes.get(ATTR_FRIENDLY_NAME)
@ -359,3 +404,133 @@ async def test_timer_restarted_event(hass):
assert results[-1].event_type == EVENT_TIMER_RESTARTED assert results[-1].event_type == EVENT_TIMER_RESTARTED
assert len(results) == 4 assert len(results) == 4
async def test_load_from_storage(hass, storage_setup):
"""Test set up from storage."""
assert await storage_setup()
state = hass.states.get(f"{DOMAIN}.timer_from_storage")
assert state.state == STATUS_IDLE
assert state.attributes.get(ATTR_FRIENDLY_NAME) == "timer from storage"
assert state.attributes.get(ATTR_EDITABLE)
async def test_editable_state_attribute(hass, storage_setup):
"""Test editable attribute."""
assert await storage_setup(config={DOMAIN: {"from_yaml": None}})
state = hass.states.get(f"{DOMAIN}.{DOMAIN}_from_storage")
assert state.state == STATUS_IDLE
assert state.attributes.get(ATTR_FRIENDLY_NAME) == "timer from storage"
assert state.attributes.get(ATTR_EDITABLE)
state = hass.states.get(f"{DOMAIN}.from_yaml")
assert not state.attributes.get(ATTR_EDITABLE)
assert state.state == STATUS_IDLE
async def test_ws_list(hass, hass_ws_client, storage_setup):
"""Test listing via WS."""
assert await storage_setup(config={DOMAIN: {"from_yaml": None}})
client = await hass_ws_client(hass)
await client.send_json({"id": 6, "type": f"{DOMAIN}/list"})
resp = await client.receive_json()
assert resp["success"]
storage_ent = "from_storage"
yaml_ent = "from_yaml"
result = {item["id"]: item for item in resp["result"]}
assert len(result) == 1
assert storage_ent in result
assert yaml_ent not in result
assert result[storage_ent][ATTR_NAME] == "timer from storage"
async def test_ws_delete(hass, hass_ws_client, storage_setup):
"""Test WS delete cleans up entity registry."""
assert await storage_setup()
timer_id = "from_storage"
timer_entity_id = f"{DOMAIN}.{DOMAIN}_{timer_id}"
ent_reg = await entity_registry.async_get_registry(hass)
state = hass.states.get(timer_entity_id)
assert state is not None
from_reg = ent_reg.async_get_entity_id(DOMAIN, DOMAIN, timer_id)
assert from_reg == timer_entity_id
client = await hass_ws_client(hass)
await client.send_json(
{"id": 6, "type": f"{DOMAIN}/delete", f"{DOMAIN}_id": f"{timer_id}"}
)
resp = await client.receive_json()
assert resp["success"]
state = hass.states.get(timer_entity_id)
assert state is None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, timer_id) is None
async def test_update(hass, hass_ws_client, storage_setup):
"""Test updating timer entity."""
assert await storage_setup()
timer_id = "from_storage"
timer_entity_id = f"{DOMAIN}.{DOMAIN}_{timer_id}"
ent_reg = await entity_registry.async_get_registry(hass)
state = hass.states.get(timer_entity_id)
assert state.attributes[ATTR_FRIENDLY_NAME] == "timer from storage"
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, timer_id) == timer_entity_id
client = await hass_ws_client(hass)
await client.send_json(
{
"id": 6,
"type": f"{DOMAIN}/update",
f"{DOMAIN}_id": f"{timer_id}",
CONF_DURATION: 33,
}
)
resp = await client.receive_json()
assert resp["success"]
state = hass.states.get(timer_entity_id)
assert state.attributes[ATTR_DURATION] == str(cv.time_period(33))
async def test_ws_create(hass, hass_ws_client, storage_setup):
"""Test create WS."""
assert await storage_setup(items=[])
timer_id = "new_timer"
timer_entity_id = f"{DOMAIN}.{timer_id}"
ent_reg = await entity_registry.async_get_registry(hass)
state = hass.states.get(timer_entity_id)
assert state is None
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, timer_id) is None
client = await hass_ws_client(hass)
await client.send_json(
{
"id": 6,
"type": f"{DOMAIN}/create",
CONF_NAME: "New Timer",
CONF_DURATION: 42,
}
)
resp = await client.receive_json()
assert resp["success"]
state = hass.states.get(timer_entity_id)
assert state.state == STATUS_IDLE
assert state.attributes[ATTR_DURATION] == str(cv.time_period(42))
assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, timer_id) == timer_entity_id