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.

homeassistant/components/timer/__init__.py

@@ -1,15 +1,26 @@
 """Support for Timers."""
 from datetime import timedelta
 import logging
+import typing
 
 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
 from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.event import async_track_point_in_utc_time
 from homeassistant.helpers.restore_state import RestoreEntity
 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
 
 _LOGGER = logging.getLogger(__name__)
@@ -17,7 +28,7 @@ _LOGGER = logging.getLogger(__name__)
 DOMAIN = "timer"
 ENTITY_ID_FORMAT = DOMAIN + ".{}"
 
-DEFAULT_DURATION = timedelta(0)
+DEFAULT_DURATION = 0
 ATTR_DURATION = "duration"
 ATTR_REMAINING = "remaining"
 CONF_DURATION = "duration"
@@ -37,6 +48,21 @@ SERVICE_PAUSE = "pause"
 SERVICE_CANCEL = "cancel"
 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):
     if value is None:
@@ -65,20 +91,55 @@ CONFIG_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)
+    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
-        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(
         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_FINISH, {}, "async_finish")
 
-    if entities:
-        await component.async_add_entities(entities)
     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):
     """Representation of a timer."""
 
-    def __init__(self, hass, object_id, name, icon, duration):
+    def __init__(self, config: typing.Dict):
         """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._duration = duration
+        self._remaining = config[CONF_DURATION]
-        self._remaining = self._duration
-        self._icon = icon
-        self._hass = hass
         self._end = 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
     def should_poll(self):
         """If entity should be polled."""
@@ -141,12 +214,12 @@ class Timer(RestoreEntity):
     @property
     def name(self):
         """Return name of the timer."""
-        return self._name
+        return self._config.get(CONF_NAME)
 
     @property
     def icon(self):
         """Return the icon to be used for this entity."""
-        return self._icon
+        return self._config.get(CONF_ICON)
 
     @property
     def state(self):
@@ -157,10 +230,16 @@ class Timer(RestoreEntity):
     def state_attributes(self):
         """Return the state attributes."""
         return {
-            ATTR_DURATION: str(self._duration),
+            ATTR_DURATION: str(self._config[CONF_DURATION]),
+            ATTR_EDITABLE: self.editable,
             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):
         """Call when entity is about to be added to Home Assistant."""
         # If not None, we got an initial value.
@@ -189,18 +268,18 @@ class Timer(RestoreEntity):
             self._end = start + self._remaining
         else:
             if newduration:
-                self._duration = newduration
+                self._config[CONF_DURATION] = newduration
                 self._remaining = newduration
             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._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):
         """Pause a timer."""
@@ -212,8 +291,8 @@ class Timer(RestoreEntity):
         self._remaining = self._end - dt_util.utcnow().replace(microsecond=0)
         self._state = STATUS_PAUSED
         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):
         """Cancel a timer."""
@@ -223,8 +302,8 @@ class Timer(RestoreEntity):
         self._state = STATUS_IDLE
         self._end = None
         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):
         """Reset and updates the states, fire finished event."""
@@ -234,8 +313,8 @@ class Timer(RestoreEntity):
         self._listener = None
         self._state = STATUS_IDLE
         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):
         """Reset and updates the states, fire finished event."""
@@ -245,5 +324,10 @@ class Timer(RestoreEntity):
         self._listener = None
         self._state = STATUS_IDLE
         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()

tests/components/timer/test_init.py

@@ -27,13 +27,17 @@ from homeassistant.components.timer import (
     STATUS_PAUSED,
 )
 from homeassistant.const import (
+    ATTR_EDITABLE,
     ATTR_FRIENDLY_NAME,
     ATTR_ICON,
+    ATTR_ID,
+    ATTR_NAME,
     CONF_ENTITY_ID,
     SERVICE_RELOAD,
 )
 from homeassistant.core import Context, CoreState
 from homeassistant.exceptions import Unauthorized
+from homeassistant.helpers import config_validation as cv, entity_registry
 from homeassistant.setup import async_setup_component
 from homeassistant.util.dt import utcnow
 
@@ -42,6 +46,38 @@ from tests.common import async_fire_time_changed
 _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):
     """Test config."""
     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 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):
@@ -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):
     """Test reload service."""
     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())
 
@@ -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_2 is not 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 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_2 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 "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 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