From 7d506bc38b73f5a81649a0fe629f054ea5dbdae6 Mon Sep 17 00:00:00 2001 From: Alexei Chetroi Date: Sat, 11 Jan 2020 11:37:57 -0500 Subject: [PATCH] Use storage collections for input_select entity management (#30604) * Refactor input_select to use _config dict. * Use collections for input_select. * Add tests. * Move async_setup to top. * Cleanup. * async_write_ha_state() * Update homeassistant/components/input_select/__init__.py Co-Authored-By: Paulus Schoutsen --- .../components/input_select/__init__.py | 167 +++++++++++---- tests/components/input_select/test_init.py | 196 +++++++++++++++++- .../input_select/test_reproduce_state.py | 2 +- 3 files changed, 322 insertions(+), 43 deletions(-) diff --git a/homeassistant/components/input_select/__init__.py b/homeassistant/components/input_select/__init__.py index b2b4b2083e8..937af76ed4f 100644 --- a/homeassistant/components/input_select/__init__.py +++ b/homeassistant/components/input_select/__init__.py @@ -1,13 +1,24 @@ """Support to select an option from a list.""" 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.restore_state import RestoreEntity import homeassistant.helpers.service +from homeassistant.helpers.storage import Store +from homeassistant.helpers.typing import ConfigType, HomeAssistantType, ServiceCallType _LOGGER = logging.getLogger(__name__) @@ -21,13 +32,24 @@ ATTR_OPTION = "option" ATTR_OPTIONS = "options" SERVICE_SELECT_OPTION = "select_option" - - SERVICE_SELECT_NEXT = "select_next" - SERVICE_SELECT_PREVIOUS = "select_previous" - SERVICE_SET_OPTIONS = "set_options" +STORAGE_KEY = DOMAIN +STORAGE_VERSION = 1 + +CREATE_FIELDS = { + vol.Required(CONF_NAME): vol.All(str, vol.Length(min=1)), + vol.Required(CONF_OPTIONS): vol.All(cv.ensure_list, vol.Length(min=1), [cv.string]), + vol.Optional(CONF_INITIAL): cv.string, + vol.Optional(CONF_ICON): cv.icon, +} +UPDATE_FIELDS = { + vol.Optional(CONF_NAME): cv.string, + vol.Optional(CONF_OPTIONS): vol.All(cv.ensure_list, vol.Length(min=1), [cv.string]), + vol.Optional(CONF_INITIAL): cv.string, + vol.Optional(CONF_ICON): cv.icon, +} def _cv_input_select(cfg): @@ -65,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 select.""" 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, InputSelect.from_yaml + ) - async def reload_service_handler(service_call): - """Remove all entities and load new ones from config.""" - conf = await component.async_prepare_reload() - if conf is None: + storage_collection = InputSelectStorageCollection( + Store(hass, STORAGE_VERSION, STORAGE_KEY), + logging.getLogger(f"{__name__}.storage_collection"), + id_manager, + ) + collection.attach_entity_component_collection( + component, storage_collection, InputSelect + ) + + 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(conf) - if new_entities: - await component.async_add_entities(new_entities) + + 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: + """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, @@ -95,11 +154,11 @@ async def async_setup(hass, config): ) component.async_register_entity_service( - SERVICE_SELECT_NEXT, {}, lambda entity, call: entity.async_offset_index(1), + SERVICE_SELECT_NEXT, {}, lambda entity, call: entity.async_offset_index(1) ) component.async_register_entity_service( - SERVICE_SELECT_PREVIOUS, {}, lambda entity, call: entity.async_offset_index(-1), + SERVICE_SELECT_PREVIOUS, {}, lambda entity, call: entity.async_offset_index(-1) ) component.async_register_entity_service( @@ -112,35 +171,46 @@ async def async_setup(hass, config): "async_set_options", ) - if entities: - await component.async_add_entities(entities) return True -async def _async_process_config(config): - """Process config and create list of entities.""" - entities = [] +class InputSelectStorageCollection(collection.StorageCollection): + """Input storage based collection.""" - for object_id, cfg in config[DOMAIN].items(): - name = cfg.get(CONF_NAME) - options = cfg.get(CONF_OPTIONS) - initial = cfg.get(CONF_INITIAL) - icon = cfg.get(CONF_ICON) - entities.append(InputSelect(object_id, name, initial, options, icon)) + CREATE_SCHEMA = vol.Schema(vol.All(CREATE_FIELDS, _cv_input_select)) + UPDATE_SCHEMA = vol.Schema(UPDATE_FIELDS) - 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 _cv_input_select({**data, **update_data}) class InputSelect(RestoreEntity): """Representation of a select input.""" - def __init__(self, object_id, name, initial, options, icon): + def __init__(self, config: typing.Dict): """Initialize a select input.""" - self.entity_id = ENTITY_ID_FORMAT.format(object_id) - self._name = name - self._current_option = initial - self._options = options - self._icon = icon + self._config = config + self.editable = True + self._current_option = config.get(CONF_INITIAL) + + @classmethod + def from_yaml(cls, config: typing.Dict) -> "InputSelect": + """Return entity instance initialized from yaml storage.""" + input_select = cls(config) + input_select.entity_id = ENTITY_ID_FORMAT.format(config[CONF_ID]) + input_select.editable = False + return input_select async def async_added_to_hass(self): """Run when entity about to be added.""" @@ -162,12 +232,17 @@ class InputSelect(RestoreEntity): @property def name(self): """Return the name of the select input.""" - 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 _options(self) -> typing.List[str]: + """Return a list of selection options.""" + return self._config[CONF_OPTIONS] @property def state(self): @@ -177,7 +252,12 @@ class InputSelect(RestoreEntity): @property def state_attributes(self): """Return the state attributes.""" - return {ATTR_OPTIONS: self._options} + return {ATTR_OPTIONS: self._config[ATTR_OPTIONS], ATTR_EDITABLE: self.editable} + + @property + def unique_id(self) -> typing.Optional[str]: + """Return unique id for the entity.""" + return self._config[CONF_ID] async def async_select_option(self, option): """Select new option.""" @@ -189,17 +269,22 @@ class InputSelect(RestoreEntity): ) return self._current_option = option - await self.async_update_ha_state() + self.async_write_ha_state() async def async_offset_index(self, offset): """Offset current index.""" current_index = self._options.index(self._current_option) new_index = (current_index + offset) % len(self._options) self._current_option = self._options[new_index] - await self.async_update_ha_state() + self.async_write_ha_state() async def async_set_options(self, options): """Set options.""" self._current_option = options[0] - self._options = options - await self.async_update_ha_state() + self._config[CONF_OPTIONS] = options + 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() diff --git a/tests/components/input_select/test_init.py b/tests/components/input_select/test_init.py index 8fda80cd3d2..13669ea507f 100644 --- a/tests/components/input_select/test_init.py +++ b/tests/components/input_select/test_init.py @@ -7,6 +7,7 @@ import pytest from homeassistant.components.input_select import ( ATTR_OPTION, ATTR_OPTIONS, + CONF_INITIAL, DOMAIN, SERVICE_SELECT_NEXT, SERVICE_SELECT_OPTION, @@ -14,19 +15,54 @@ from homeassistant.components.input_select import ( SERVICE_SET_OPTIONS, ) from homeassistant.const import ( + ATTR_EDITABLE, ATTR_ENTITY_ID, ATTR_FRIENDLY_NAME, ATTR_ICON, + ATTR_NAME, SERVICE_RELOAD, ) from homeassistant.core import Context, State from homeassistant.exceptions import Unauthorized +from homeassistant.helpers import entity_registry from homeassistant.loader import bind_hass from homeassistant.setup import async_setup_component from tests.common import mock_restore_cache +@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", + "options": ["storage option 1", "storage option 2"], + } + ] + }, + } + 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 + + @bind_hass def select_option(hass, entity_id, option): """Set value of input_select. @@ -329,6 +365,7 @@ async def test_input_select_context(hass, hass_admin_user): async def test_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) assert await async_setup_component( hass, @@ -358,6 +395,9 @@ async def test_reload(hass, hass_admin_user, hass_read_only_user): assert state_3 is None assert "middle option" == state_1.state assert "an option" == 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", @@ -400,5 +440,159 @@ async def test_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 "reloaded option" == state_2.state + assert "an option" == state_2.state assert "newer option" == state_3.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 + + +async def test_load_from_storage(hass, storage_setup): + """Test set up from storage.""" + assert await storage_setup() + state = hass.states.get(f"{DOMAIN}.from_storage") + assert state.state == "storage option 1" + 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": {"options": ["yaml option", "other option"]}}} + ) + + state = hass.states.get(f"{DOMAIN}.from_storage") + assert state.state == "storage option 1" + 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 == "yaml option" + 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": {"options": ["yaml option"]}}} + ) + + 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 + + +async def test_update(hass, hass_ws_client, storage_setup): + """Test updating min/max updates the state.""" + + items = [ + { + "id": "from_storage", + "name": "from storage", + "options": ["yaml update 1", "yaml update 2"], + } + ] + assert await storage_setup(items) + + 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.attributes[ATTR_OPTIONS] == ["yaml update 1", "yaml update 2"] + 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}/update", + f"{DOMAIN}_id": f"{input_id}", + "options": ["new option", "newer option"], + CONF_INITIAL: "newer option", + } + ) + resp = await client.receive_json() + assert resp["success"] + + state = hass.states.get(input_entity_id) + assert state.attributes[ATTR_OPTIONS] == ["new option", "newer option"] + + await client.send_json( + { + "id": 7, + "type": f"{DOMAIN}/update", + f"{DOMAIN}_id": f"{input_id}", + "options": ["new option", "no newer option"], + } + ) + resp = await client.receive_json() + assert not resp["success"] + + +async def test_ws_create(hass, hass_ws_client, storage_setup): + """Test create WS.""" + assert await storage_setup(items=[]) + + input_id = "new_input" + 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", + "name": "New Input", + "options": ["new option", "even newer option"], + "initial": "even newer option", + } + ) + resp = await client.receive_json() + assert resp["success"] + + state = hass.states.get(input_entity_id) + assert state.state == "even newer option" diff --git a/tests/components/input_select/test_reproduce_state.py b/tests/components/input_select/test_reproduce_state.py index 469c258cb4b..ed1f9f45e43 100644 --- a/tests/components/input_select/test_reproduce_state.py +++ b/tests/components/input_select/test_reproduce_state.py @@ -68,5 +68,5 @@ async def test_reproducing_states(hass, caplog): ) # These should fail if options weren't changed to VALID_OPTION_SET2 - assert hass.states.get(ENTITY).attributes == {"options": VALID_OPTION_SET2} + assert hass.states.get(ENTITY).attributes["options"] == VALID_OPTION_SET2 assert hass.states.get(ENTITY).state == VALID_OPTION5