Add remove commmand to entity registry (#20597)

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'], [{
+        msg['id'], [_entry_dict(entry) for entry in registry.entities.values()]
-            '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()]
     ))
 
 
@@ -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,
     }

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):

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