Use collection helpers for input_boolean (#30514)

* Refactor input_boolean to use config dict for instantiation.

* Refactor input_boolean to use YamlCollection.

* Add storage collection to input_boolean.

* Update homeassistant/components/input_boolean/__init__.py

Co-Authored-By: Paulus Schoutsen <paulus@home-assistant.io>

* Doh.

* Address comments.

* Add editable device state attribute.

* Clean up entities from entity registry on removal.

Reload yaml from correct source.

* Add tests.

* Update homeassistant/components/input_boolean/__init__.py

Co-Authored-By: Paulus Schoutsen <paulus@home-assistant.io>

* Don't reset entity state on updates.

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/input_boolean/__init__.py

@@ -1,10 +1,13 @@
 """Support to keep track of user controlled booleans for within automation."""
 import logging
+import typing
 
 import voluptuous as vol
 
 from homeassistant.const import (
+    ATTR_EDITABLE,
     CONF_ICON,
+    CONF_ID,
     CONF_NAME,
     SERVICE_RELOAD,
     SERVICE_TOGGLE,
@@ -12,11 +15,15 @@ from homeassistant.const import (
     SERVICE_TURN_ON,
     STATE_ON,
 )
+from homeassistant.core import callback
+from homeassistant.helpers import collection, entity_registry
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import ToggleEntity
 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.loader import bind_hass
 
 DOMAIN = "input_boolean"
@@ -27,23 +34,47 @@ _LOGGER = logging.getLogger(__name__)
 
 CONF_INITIAL = "initial"
 
-CONFIG_SCHEMA = vol.Schema(
+CREATE_FIELDS = {
-    {
+    vol.Required(CONF_NAME): vol.All(str, vol.Length(min=1)),
-        DOMAIN: cv.schema_with_slug_keys(
+    vol.Optional(CONF_INITIAL): cv.boolean,
-            vol.Any(
+    vol.Optional(CONF_ICON): cv.icon,
-                {
+}
+
+UPDATE_FIELDS = {
     vol.Optional(CONF_NAME): cv.string,
     vol.Optional(CONF_INITIAL): cv.boolean,
     vol.Optional(CONF_ICON): cv.icon,
-                },
+}
-                None,
+
-            )
+CONFIG_SCHEMA = vol.Schema(
-        )
+    {DOMAIN: cv.schema_with_slug_keys(vol.Any(UPDATE_FIELDS, None))},
-    },
     extra=vol.ALLOW_EXTRA,
 )
 
 RELOAD_SERVICE_SCHEMA = vol.Schema({})
+STORAGE_KEY = DOMAIN
+STORAGE_VERSION = 1
+
+
+class InputBooleanStorageCollection(collection.StorageCollection):
+    """Input boolean collection stored in storage."""
+
+    CREATE_SCHEMA = vol.Schema(CREATE_FIELDS)
+    UPDATE_SCHEMA = vol.Schema(UPDATE_FIELDS)
+
+    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 {**data, **update_data}
 
 
 @bind_hass
@@ -52,20 +83,57 @@ def is_on(hass, entity_id):
     return hass.states.is_state(entity_id, STATE_ON)
 
 
-async def async_setup(hass, config):
+async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     """Set up an input boolean."""
     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, lambda conf: InputBoolean(conf, from_yaml=True)
+    )
 
-    async def reload_service_handler(service_call):
+    storage_collection = InputBooleanStorageCollection(
+        Store(hass, STORAGE_VERSION, STORAGE_KEY),
+        logging.getLogger(f"{__name__}_storage_collection"),
+        id_manager,
+    )
+    collection.attach_entity_component_collection(
+        component, storage_collection, InputBoolean
+    )
+
+    await yaml_collection.async_load(
+        [{CONF_ID: id_, **(conf or {})} for id_, conf 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
+
+        ent_reg = await entity_registry.async_get_registry(hass)
+        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:
         """Remove all input booleans and load new ones from config."""
-        conf = await component.async_prepare_reload()
+        conf = await component.async_prepare_reload(skip_reset=True)
         if conf is None:
             return
-        new_entities = await _async_process_config(conf)
+        await yaml_collection.async_load(
-        if new_entities:
+            [{CONF_ID: id_, **(conf or {})} for id_, conf in conf[DOMAIN].items()]
-            await component.async_add_entities(new_entities)
+        )
 
     homeassistant.helpers.service.async_register_admin_service(
         hass,
@@ -81,38 +149,20 @@ async def async_setup(hass, config):
 
     component.async_register_entity_service(SERVICE_TOGGLE, {}, "async_toggle")
 
-    if entities:
-        await component.async_add_entities(entities)
-
     return True
 
 
-async def _async_process_config(config):
-    """Process config and create list of entities."""
-    entities = []
-
-    for object_id, cfg in config[DOMAIN].items():
-        if not cfg:
-            cfg = {}
-
-        name = cfg.get(CONF_NAME)
-        initial = cfg.get(CONF_INITIAL)
-        icon = cfg.get(CONF_ICON)
-
-        entities.append(InputBoolean(object_id, name, initial, icon))
-
-    return entities
-
-
 class InputBoolean(ToggleEntity, RestoreEntity):
     """Representation of a boolean input."""
 
-    def __init__(self, object_id, name, initial, icon):
+    def __init__(self, config: typing.Optional[dict], from_yaml: bool = False):
         """Initialize a boolean input."""
-        self.entity_id = ENTITY_ID_FORMAT.format(object_id)
+        self._config = config
-        self._name = name
+        self._editable = True
-        self._state = initial
+        self._state = config.get(CONF_INITIAL)
-        self._icon = icon
+        if from_yaml:
+            self._editable = False
+            self.entity_id = ENTITY_ID_FORMAT.format(self.unique_id)
 
     @property
     def should_poll(self):
@@ -122,18 +172,28 @@ class InputBoolean(ToggleEntity, RestoreEntity):
     @property
     def name(self):
         """Return name of the boolean input."""
-        return self._name
+        return self._config.get(CONF_NAME)
+
+    @property
+    def state_attributes(self):
+        """Return the state attributes of the entity."""
+        return {ATTR_EDITABLE: self._editable}
 
     @property
     def icon(self):
         """Return the icon to be used for this entity."""
-        return self._icon
+        return self._config.get(CONF_ICON)
 
     @property
     def is_on(self):
         """Return true if entity is on."""
         return self._state
 
+    @property
+    def unique_id(self):
+        """Return a unique ID for the person."""
+        return self._config[CONF_ID]
+
     async def async_added_to_hass(self):
         """Call when entity about to be added to hass."""
         # If not None, we got an initial value.
@@ -153,3 +213,8 @@ class InputBoolean(ToggleEntity, RestoreEntity):
         """Turn the entity off."""
         self._state = False
         await self.async_update_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()

homeassistant/const.py

@@ -318,6 +318,7 @@ ATTR_GPS_ACCURACY = "gps_accuracy"
 ATTR_ASSUMED_STATE = "assumed_state"
 ATTR_STATE = "state"
 
+ATTR_EDITABLE = "editable"
 ATTR_OPTION = "option"
 
 # Bitfield of supported component features for the entity

tests/components/input_boolean/test_init.py

@@ -3,11 +3,15 @@
 import logging
 from unittest.mock import patch
 
+import pytest
+
 from homeassistant.components.input_boolean import CONF_INITIAL, DOMAIN, is_on
 from homeassistant.const import (
+    ATTR_EDITABLE,
     ATTR_ENTITY_ID,
     ATTR_FRIENDLY_NAME,
     ATTR_ICON,
+    ATTR_NAME,
     SERVICE_RELOAD,
     SERVICE_TOGGLE,
     SERVICE_TURN_OFF,
@@ -16,6 +20,7 @@ from homeassistant.const import (
     STATE_ON,
 )
 from homeassistant.core import Context, CoreState, State
+from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
 
 from tests.common import mock_component, mock_restore_cache
@@ -23,6 +28,26 @@ from tests.common import mock_component, mock_restore_cache
 _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": [{"id": "from_storage", "name": "from storage"}]},
+            }
+        else:
+            hass_storage[DOMAIN] = 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}]
@@ -169,6 +194,7 @@ async def test_input_boolean_context(hass, hass_admin_user):
 async def test_reload(hass, hass_admin_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())
 
@@ -196,6 +222,10 @@ async def test_reload(hass, hass_admin_user):
     assert state_3 is None
     assert STATE_ON == state_2.state
 
+    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
+
     with patch(
         "homeassistant.config.load_yaml_config_file",
         autospec=True,
@@ -229,6 +259,78 @@ async def test_reload(hass, hass_admin_user):
     assert state_2 is not None
     assert state_3 is not None
 
-    assert STATE_OFF == state_2.state
+    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 STATE_ON == state_2.state  # reload is not supposed to change entity state
     assert "Hello World reloaded" == state_2.attributes.get(ATTR_FRIENDLY_NAME)
     assert "mdi:work_reloaded" == state_2.attributes.get(ATTR_ICON)
+
+
+async def test_load_person_storage(hass, storage_setup):
+    """Test set up from storage."""
+    assert await storage_setup()
+    state = hass.states.get(f"{DOMAIN}.from_storage")
+    assert state.state == STATE_OFF
+    assert state.attributes.get(ATTR_FRIENDLY_NAME) == "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}.from_storage")
+    assert state.state == STATE_OFF
+    assert state.attributes.get(ATTR_FRIENDLY_NAME) == "from storage"
+    assert state.attributes.get(ATTR_EDITABLE)
+
+    state = hass.states.get(f"{DOMAIN}.from_yaml")
+    assert state.state == STATE_OFF
+    assert not state.attributes.get(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": 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] == "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}.{input_id}"
+    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) is not None
+
+    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