From 6740249befc07f7dc683ad634e05b1dc94f49093 Mon Sep 17 00:00:00 2001 From: Alexei Chetroi Date: Tue, 7 Jan 2020 02:53:43 -0500 Subject: [PATCH] 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 * 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 * Don't reset entity state on updates. Co-authored-by: Paulus Schoutsen --- .../components/input_boolean/__init__.py | 157 +++++++++++++----- homeassistant/const.py | 1 + tests/components/input_boolean/test_init.py | 104 +++++++++++- 3 files changed, 215 insertions(+), 47 deletions(-) diff --git a/homeassistant/components/input_boolean/__init__.py b/homeassistant/components/input_boolean/__init__.py index 326244a7e4c..a12c6552399 100644 --- a/homeassistant/components/input_boolean/__init__.py +++ b/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" +CREATE_FIELDS = { + vol.Required(CONF_NAME): vol.All(str, vol.Length(min=1)), + vol.Optional(CONF_INITIAL): cv.boolean, + 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, +} + CONFIG_SCHEMA = vol.Schema( - { - DOMAIN: cv.schema_with_slug_keys( - vol.Any( - { - vol.Optional(CONF_NAME): cv.string, - vol.Optional(CONF_INITIAL): cv.boolean, - vol.Optional(CONF_ICON): cv.icon, - }, - None, - ) - ) - }, + {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) - if new_entities: - await component.async_add_entities(new_entities) + await yaml_collection.async_load( + [{CONF_ID: id_, **(conf or {})} for id_, conf in conf[DOMAIN].items()] + ) 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._name = name - self._state = initial - self._icon = icon + self._config = config + self._editable = True + self._state = config.get(CONF_INITIAL) + 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() diff --git a/homeassistant/const.py b/homeassistant/const.py index ae055c52f10..a9642e88e15 100644 --- a/homeassistant/const.py +++ b/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 diff --git a/tests/components/input_boolean/test_init.py b/tests/components/input_boolean/test_init.py index 2d504114c78..1ab7e9a3d13 100644 --- a/tests/components/input_boolean/test_init.py +++ b/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