Use collection helpers for input_datetime component (#30815)

* Refactor input_datetime.
Keep the state as datetime, but format accordingly to has_time and
has_date.

* Use config dict for input_datetime.
* Add tests.
* Lint
Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/input_datetime/__init__.py

@@ -1,20 +1,27 @@
 """Support to select a date and/or a time."""
 import datetime
 import logging
+import typing
 
 import voluptuous as vol
 
 from homeassistant.const import (
     ATTR_DATE,
+    ATTR_EDITABLE,
     ATTR_TIME,
     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.restore_state import RestoreEntity
 import homeassistant.helpers.service
+from homeassistant.helpers.storage import Store
+from homeassistant.helpers.typing import ConfigType, HomeAssistantType, ServiceCallType
 from homeassistant.util import dt as dt_util
 
 _LOGGER = logging.getLogger(__name__)
@@ -27,10 +34,29 @@ CONF_HAS_TIME = "has_time"
 CONF_INITIAL = "initial"
 
 DEFAULT_VALUE = "1970-01-01 00:00:00"
+DEFAULT_DATE = datetime.date(1970, 1, 1)
+DEFAULT_TIME = datetime.time(0, 0, 0)
 
 ATTR_DATETIME = "datetime"
 
 SERVICE_SET_DATETIME = "set_datetime"
+STORAGE_KEY = DOMAIN
+STORAGE_VERSION = 1
+
+CREATE_FIELDS = {
+    vol.Required(CONF_NAME): vol.All(str, vol.Length(min=1)),
+    vol.Optional(CONF_HAS_DATE, default=False): cv.boolean,
+    vol.Optional(CONF_HAS_TIME, default=False): cv.boolean,
+    vol.Optional(CONF_ICON): cv.icon,
+    vol.Optional(CONF_INITIAL): cv.string,
+}
+UPDATE_FIELDS = {
+    vol.Optional(CONF_NAME): cv.string,
+    vol.Optional(CONF_HAS_DATE): cv.boolean,
+    vol.Optional(CONF_HAS_TIME): cv.boolean,
+    vol.Optional(CONF_ICON): cv.icon,
+    vol.Optional(CONF_INITIAL): cv.string,
+}
 
 
 def has_date_or_time(conf):
@@ -61,20 +87,57 @@ 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 an input datetime."""
     component = EntityComponent(_LOGGER, DOMAIN, hass)
+    id_manager = collection.IDManager()
 
-    entities = await _async_process_config(config)
+    yaml_collection = collection.YamlCollection(
+        logging.getLogger(f"{__name__}.yaml_collection"), id_manager
+    )
+    collection.attach_entity_component_collection(
+        component, yaml_collection, InputDatetime.from_yaml
+    )
 
-    async def reload_service_handler(service_call):
+    storage_collection = DateTimeStorageCollection(
-        """Remove all entities 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, InputDatetime
+    )
+
+    await yaml_collection.async_load(
+        [{CONF_ID: id_, **cfg} for id_, cfg in config.get(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(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.get(DOMAIN, {}).items()]
+        )
 
     homeassistant.helpers.service.async_register_admin_service(
         hass,
@@ -119,68 +182,79 @@ async def async_setup(hass, config):
         async_set_datetime_service,
     )
 
-    if entities:
-        await component.async_add_entities(entities)
     return True
 
 
-async def _async_process_config(config):
+class DateTimeStorageCollection(collection.StorageCollection):
-    """Process config and create list of entities."""
+    """Input storage based collection."""
-    entities = []
 
-    for object_id, cfg in config[DOMAIN].items():
+    CREATE_SCHEMA = vol.Schema(vol.All(CREATE_FIELDS, has_date_or_time))
-        name = cfg.get(CONF_NAME)
+    UPDATE_SCHEMA = vol.Schema(UPDATE_FIELDS)
-        has_time = cfg.get(CONF_HAS_TIME)
-        has_date = cfg.get(CONF_HAS_DATE)
-        icon = cfg.get(CONF_ICON)
-        initial = cfg.get(CONF_INITIAL)
-        entities.append(
-            InputDatetime(object_id, name, has_date, has_time, icon, initial)
-        )
 
-    return entities
+    async def _process_create_data(self, data: typing.Dict) -> typing.Dict:
+        """Validate the config is valid."""
+        return self.CREATE_SCHEMA(data)
+
+    @callback
+    def _get_suggested_id(self, info: typing.Dict) -> str:
+        """Suggest an ID based on the config."""
+        return info[CONF_NAME]
+
+    async def _update_data(self, data: dict, update_data: typing.Dict) -> typing.Dict:
+        """Return a new updated data object."""
+        update_data = self.UPDATE_SCHEMA(update_data)
+        return has_date_or_time({**data, **update_data})
 
 
 class InputDatetime(RestoreEntity):
     """Representation of a datetime input."""
 
-    def __init__(self, object_id, name, has_date, has_time, icon, initial):
+    def __init__(self, config: typing.Dict) -> None:
         """Initialize a select input."""
-        self.entity_id = ENTITY_ID_FORMAT.format(object_id)
+        self._config = config
-        self._name = name
+        self.editable = True
-        self.has_date = has_date
-        self.has_time = has_time
-        self._icon = icon
-        self._initial = initial
         self._current_datetime = None
+        initial = config.get(CONF_INITIAL)
+        if initial:
+            if self.has_date and self.has_time:
+                self._current_datetime = dt_util.parse_datetime(initial)
+            elif self.has_date:
+                date = dt_util.parse_date(initial)
+                self._current_datetime = datetime.datetime.combine(date, DEFAULT_TIME)
+            else:
+                time = dt_util.parse_time(initial)
+                self._current_datetime = datetime.datetime.combine(DEFAULT_DATE, time)
+
+    @classmethod
+    def from_yaml(cls, config: typing.Dict) -> "InputDatetime":
+        """Return entity instance initialized from yaml storage."""
+        input_dt = cls(config)
+        input_dt.entity_id = ENTITY_ID_FORMAT.format(config[CONF_ID])
+        input_dt.editable = False
+        return input_dt
 
     async def async_added_to_hass(self):
         """Run when entity about to be added."""
         await super().async_added_to_hass()
-        restore_val = None
 
-        # Priority 1: Initial State
+        # Priority 1: Initial value
-        if self._initial is not None:
+        if self.state is not None:
-            restore_val = self._initial
+            return
 
         # Priority 2: Old state
-        if restore_val is None:
+        old_state = await self.async_get_last_state()
-            old_state = await self.async_get_last_state()
+        if old_state is None:
-            if old_state is not None:
+            self._current_datetime = dt_util.parse_datetime(DEFAULT_VALUE)
-                restore_val = old_state.state
+            return
 
-        if not self.has_date:
+        if self.has_date and self.has_time:
-            if not restore_val:
+            self._current_datetime = dt_util.parse_datetime(old_state.state)
-                restore_val = DEFAULT_VALUE.split()[1]
+        elif self.has_date:
-            self._current_datetime = dt_util.parse_time(restore_val)
+            date = dt_util.parse_date(old_state.state)
-        elif not self.has_time:
+            self._current_datetime = datetime.datetime.combine(date, DEFAULT_TIME)
-            if not restore_val:
-                restore_val = DEFAULT_VALUE.split()[0]
-            self._current_datetime = dt_util.parse_date(restore_val)
         else:
-            if not restore_val:
+            time = dt_util.parse_time(old_state.state)
-                restore_val = DEFAULT_VALUE
+            self._current_datetime = datetime.datetime.combine(DEFAULT_DATE, time)
-            self._current_datetime = dt_util.parse_datetime(restore_val)
 
     @property
     def should_poll(self):
@@ -190,22 +264,43 @@ class InputDatetime(RestoreEntity):
     @property
     def name(self):
         """Return the name of the select input."""
-        return self._name
+        return self._config.get(CONF_NAME)
+
+    @property
+    def has_date(self) -> bool:
+        """Return True if entity has date."""
+        return self._config[CONF_HAS_DATE]
+
+    @property
+    def has_time(self) -> bool:
+        """Return True if entity has time."""
+        return self._config[CONF_HAS_TIME]
 
     @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):
         """Return the state of the component."""
-        return self._current_datetime
+        if self._current_datetime is None:
+            return None
+
+        if self.has_date and self.has_time:
+            return self._current_datetime
+        if self.has_date:
+            return self._current_datetime.date()
+        return self._current_datetime.time()
 
     @property
     def state_attributes(self):
         """Return the state attributes."""
-        attrs = {"has_date": self.has_date, "has_time": self.has_time}
+        attrs = {
+            ATTR_EDITABLE: self.editable,
+            CONF_HAS_DATE: self.has_date,
+            CONF_HAS_TIME: self.has_time,
+        }
 
         if self._current_datetime is None:
             return attrs
@@ -236,13 +331,27 @@ class InputDatetime(RestoreEntity):
 
         return attrs
 
+    @property
+    def unique_id(self) -> typing.Optional[str]:
+        """Return unique id of the entity."""
+        return self._config[CONF_ID]
+
     def async_set_datetime(self, date_val, time_val):
         """Set a new date / time."""
         if self.has_date and self.has_time and date_val and time_val:
             self._current_datetime = datetime.datetime.combine(date_val, time_val)
         elif self.has_date and not self.has_time and date_val:
-            self._current_datetime = date_val
+            self._current_datetime = datetime.datetime.combine(
+                date_val, self._current_datetime.time()
+            )
         if self.has_time and not self.has_date and time_val:
-            self._current_datetime = time_val
+            self._current_datetime = datetime.datetime.combine(
+                self._current_datetime.date(), time_val
+            )
 
-        self.async_schedule_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/input_datetime/test_init.py

@@ -9,18 +9,64 @@ import voluptuous as vol
 from homeassistant.components.input_datetime import (
     ATTR_DATE,
     ATTR_DATETIME,
+    ATTR_EDITABLE,
     ATTR_TIME,
+    CONF_HAS_DATE,
+    CONF_HAS_TIME,
+    CONF_ID,
+    CONF_INITIAL,
+    CONF_NAME,
+    DEFAULT_TIME,
     DOMAIN,
     SERVICE_RELOAD,
     SERVICE_SET_DATETIME,
 )
-from homeassistant.const import ATTR_ENTITY_ID
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_FRIENDLY_NAME, ATTR_NAME
 from homeassistant.core import Context, CoreState, State
 from homeassistant.exceptions import Unauthorized
+from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
 
 from tests.common import mock_restore_cache
 
+INITIAL_DATE = "2020-01-10"
+INITIAL_TIME = "23:45:56"
+INITIAL_DATETIME = f"{INITIAL_DATE} {INITIAL_TIME}"
+
+
+@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": [
+                        {
+                            CONF_ID: "from_storage",
+                            CONF_NAME: "datetime from storage",
+                            CONF_INITIAL: INITIAL_DATETIME,
+                            CONF_HAS_DATE: True,
+                            CONF_HAS_TIME: True,
+                        }
+                    ]
+                },
+            }
+        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 async_set_date_and_time(hass, entity_id, dt_value):
     """Set date and / or time of input_datetime."""
@@ -318,11 +364,7 @@ async def test_reload(hass, hass_admin_user, hass_read_only_user):
     assert await async_setup_component(
         hass,
         DOMAIN,
-        {
+        {DOMAIN: {"dt1": {"has_time": False, "has_date": True, "initial": "2019-1-1"}}},
-            DOMAIN: {
-                "dt1": {"has_time": False, "has_date": True, "initial": "2019-1-1"},
-            }
-        },
     )
 
     assert count_start + 1 == len(hass.states.async_entity_ids())
@@ -365,8 +407,169 @@ async def test_reload(hass, hass_admin_user, hass_read_only_user):
     state_1 = hass.states.get("input_datetime.dt1")
     state_2 = hass.states.get("input_datetime.dt2")
 
-    dt_obj = datetime.datetime(2019, 1, 1, 23, 32)
     assert state_1 is not None
     assert state_2 is not None
-    assert str(dt_obj.time()) == state_1.state
+    assert str(DEFAULT_TIME) == state_1.state
     assert str(datetime.datetime(1970, 1, 1, 0, 0)) == state_2.state
+
+
+async def test_load_from_storage(hass, storage_setup):
+    """Test set up from storage."""
+    assert await storage_setup()
+    state = hass.states.get(f"{DOMAIN}.datetime_from_storage")
+    assert state.state == INITIAL_DATETIME
+    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": {
+                    CONF_HAS_DATE: True,
+                    CONF_HAS_TIME: True,
+                    CONF_NAME: "yaml datetime",
+                    CONF_INITIAL: "2001-01-02 12:34:56",
+                }
+            }
+        }
+    )
+
+    state = hass.states.get(f"{DOMAIN}.datetime_from_storage")
+    assert state.state == INITIAL_DATETIME
+    assert state.attributes.get(ATTR_EDITABLE)
+
+    state = hass.states.get(f"{DOMAIN}.from_yaml")
+    assert state.state == "2001-01-02 12:34:56"
+    assert not state.attributes[ATTR_EDITABLE]
+
+
+async def test_ws_list(hass, hass_ws_client, storage_setup):
+    """Test listing via WS."""
+    assert await storage_setup(config={DOMAIN: {"from_yaml": {CONF_HAS_DATE: True}}})
+
+    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] == "datetime from storage"
+
+
+async def test_ws_delete(hass, hass_ws_client, storage_setup):
+    """Test WS delete cleans up entity registry."""
+    assert await storage_setup()
+
+    input_id = "from_storage"
+    input_entity_id = f"{DOMAIN}.datetime_from_storage"
+    ent_reg = await entity_registry.async_get_registry(hass)
+
+    state = hass.states.get(input_entity_id)
+    assert state is not None
+    assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, input_id) == input_entity_id
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json(
+        {"id": 6, "type": f"{DOMAIN}/delete", f"{DOMAIN}_id": f"{input_id}"}
+    )
+    resp = await client.receive_json()
+    assert resp["success"]
+
+    state = hass.states.get(input_entity_id)
+    assert state is None
+    assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, input_id) is None
+
+
+async def test_update(hass, hass_ws_client, storage_setup):
+    """Test updating min/max updates the state."""
+
+    assert await storage_setup()
+
+    input_id = "from_storage"
+    input_entity_id = f"{DOMAIN}.datetime_from_storage"
+    ent_reg = await entity_registry.async_get_registry(hass)
+
+    state = hass.states.get(input_entity_id)
+    assert state.attributes[ATTR_FRIENDLY_NAME] == "datetime from storage"
+    assert state.state == INITIAL_DATETIME
+    assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, input_id) == input_entity_id
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json(
+        {
+            "id": 6,
+            "type": f"{DOMAIN}/update",
+            f"{DOMAIN}_id": f"{input_id}",
+            ATTR_NAME: "even newer name",
+            CONF_HAS_DATE: False,
+        }
+    )
+    resp = await client.receive_json()
+    assert resp["success"]
+
+    state = hass.states.get(input_entity_id)
+    assert state.state == INITIAL_TIME
+    assert state.attributes[ATTR_FRIENDLY_NAME] == "even newer name"
+
+
+async def test_ws_create(hass, hass_ws_client, storage_setup):
+    """Test create WS."""
+    assert await storage_setup(items=[])
+
+    input_id = "new_datetime"
+    input_entity_id = f"{DOMAIN}.{input_id}"
+    ent_reg = await entity_registry.async_get_registry(hass)
+
+    state = hass.states.get(input_entity_id)
+    assert state is None
+    assert ent_reg.async_get_entity_id(DOMAIN, DOMAIN, input_id) is None
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json(
+        {
+            "id": 6,
+            "type": f"{DOMAIN}/create",
+            CONF_NAME: "New DateTime",
+            CONF_INITIAL: "1991-01-02 01:02:03",
+            CONF_HAS_DATE: True,
+            CONF_HAS_TIME: True,
+        }
+    )
+    resp = await client.receive_json()
+    assert resp["success"]
+
+    state = hass.states.get(input_entity_id)
+    assert state.state == "1991-01-02 01:02:03"
+    assert state.attributes[ATTR_FRIENDLY_NAME] == "New DateTime"
+    assert state.attributes[ATTR_EDITABLE]
+
+
+async def test_setup_no_config(hass, hass_admin_user):
+    """Test component setup with no config."""
+    count_start = len(hass.states.async_entity_ids())
+    assert await async_setup_component(hass, DOMAIN, {})
+
+    with patch(
+        "homeassistant.config.load_yaml_config_file", autospec=True, return_value={}
+    ):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_RELOAD,
+            blocking=True,
+            context=Context(user_id=hass_admin_user.id),
+        )
+        await hass.async_block_till_done()
+
+    assert count_start == len(hass.states.async_entity_ids())