diff --git a/homeassistant/components/config/entity_registry.py b/homeassistant/components/config/entity_registry.py index fdac1ad95da..39dd622540d 100644 --- a/homeassistant/components/config/entity_registry.py +++ b/homeassistant/components/config/entity_registry.py @@ -5,7 +5,8 @@ from homeassistant.core import callback from homeassistant.helpers.entity_registry import async_get_registry from homeassistant.components import websocket_api from homeassistant.components.websocket_api.const import ERR_NOT_FOUND -from homeassistant.components.websocket_api.decorators import async_response +from homeassistant.components.websocket_api.decorators import ( + async_response, require_admin) from homeassistant.helpers import config_validation as cv DEPENDENCIES = ['websocket_api'] @@ -30,6 +31,12 @@ SCHEMA_WS_UPDATE = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({ vol.Optional('new_entity_id'): str, }) +WS_TYPE_REMOVE = 'config/entity_registry/remove' +SCHEMA_WS_REMOVE = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({ + vol.Required('type'): WS_TYPE_REMOVE, + vol.Required('entity_id'): cv.entity_id +}) + async def async_setup(hass): """Enable the Entity Registry views.""" @@ -45,6 +52,10 @@ async def async_setup(hass): WS_TYPE_UPDATE, websocket_update_entity, SCHEMA_WS_UPDATE ) + hass.components.websocket_api.async_register_command( + WS_TYPE_REMOVE, websocket_remove_entity, + SCHEMA_WS_REMOVE + ) return True @@ -56,14 +67,7 @@ async def websocket_list_entities(hass, connection, msg): """ registry = await async_get_registry(hass) connection.send_message(websocket_api.result_message( - msg['id'], [{ - 'config_entry_id': entry.config_entry_id, - 'device_id': entry.device_id, - 'disabled_by': entry.disabled_by, - 'entity_id': entry.entity_id, - 'name': entry.name, - 'platform': entry.platform, - } for entry in registry.entities.values()] + msg['id'], [_entry_dict(entry) for entry in registry.entities.values()] )) @@ -86,6 +90,7 @@ async def websocket_get_entity(hass, connection, msg): )) +@require_admin @async_response async def websocket_update_entity(hass, connection, msg): """Handle update entity websocket command. @@ -125,10 +130,32 @@ async def websocket_update_entity(hass, connection, msg): )) +@require_admin +@async_response +async def websocket_remove_entity(hass, connection, msg): + """Handle remove entity websocket command. + + Async friendly. + """ + registry = await async_get_registry(hass) + + if msg['entity_id'] not in registry.entities: + connection.send_message(websocket_api.error_message( + msg['id'], ERR_NOT_FOUND, 'Entity not found')) + return + + registry.async_remove(msg['entity_id']) + connection.send_message(websocket_api.result_message(msg['id'])) + + @callback def _entry_dict(entry): """Convert entry to API format.""" return { + 'config_entry_id': entry.config_entry_id, + 'device_id': entry.device_id, + 'disabled_by': entry.disabled_by, 'entity_id': entry.entity_id, - 'name': entry.name + 'name': entry.name, + 'platform': entry.platform, } diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index e4f266854ef..6ee32f642bc 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -149,6 +149,12 @@ class EntityRegistry: self.async_schedule_save() return entity + @callback + def async_remove(self, entity_id): + """Remove an entity from registry.""" + self.entities.pop(entity_id) + self.async_schedule_save() + @callback def async_update_entity(self, entity_id, *, name=_UNDEF, new_entity_id=_UNDEF): diff --git a/tests/components/config/test_entity_registry.py b/tests/components/config/test_entity_registry.py index cd74faf1843..26903bb256b 100644 --- a/tests/components/config/test_entity_registry.py +++ b/tests/components/config/test_entity_registry.py @@ -82,6 +82,10 @@ async def test_get_entity(hass, client): msg = await client.receive_json() assert msg['result'] == { + 'config_entry_id': None, + 'device_id': None, + 'disabled_by': None, + 'platform': 'test_platform', 'entity_id': 'test_domain.name', 'name': 'Hello World' } @@ -94,6 +98,10 @@ async def test_get_entity(hass, client): msg = await client.receive_json() assert msg['result'] == { + 'config_entry_id': None, + 'device_id': None, + 'disabled_by': None, + 'platform': 'test_platform', 'entity_id': 'test_domain.no_name', 'name': None } @@ -128,6 +136,10 @@ async def test_update_entity_name(hass, client): msg = await client.receive_json() assert msg['result'] == { + 'config_entry_id': None, + 'device_id': None, + 'disabled_by': None, + 'platform': 'test_platform', 'entity_id': 'test_domain.world', 'name': 'after update' } @@ -165,6 +177,10 @@ async def test_update_entity_no_changes(hass, client): msg = await client.receive_json() assert msg['result'] == { + 'config_entry_id': None, + 'device_id': None, + 'disabled_by': None, + 'platform': 'test_platform', 'entity_id': 'test_domain.world', 'name': 'name of entity' } @@ -224,9 +240,37 @@ async def test_update_entity_id(hass, client): msg = await client.receive_json() assert msg['result'] == { + 'config_entry_id': None, + 'device_id': None, + 'disabled_by': None, + 'platform': 'test_platform', 'entity_id': 'test_domain.planet', 'name': None } assert hass.states.get('test_domain.world') is None assert hass.states.get('test_domain.planet') is not None + + +async def test_remove_entity(hass, client): + """Test removing entity.""" + registry = mock_registry(hass, { + 'test_domain.world': RegistryEntry( + entity_id='test_domain.world', + unique_id='1234', + # Using component.async_add_entities is equal to platform "domain" + platform='test_platform', + name='before update' + ) + }) + + await client.send_json({ + 'id': 6, + 'type': 'config/entity_registry/remove', + 'entity_id': 'test_domain.world', + }) + + msg = await client.receive_json() + + assert msg['success'] + assert len(registry.entities) == 0