Add lovelace websocket get and set card (#17600)

* Add ws get, set card

* lint+fix test

* Add test for set

* Added more tests, catch unsupported yaml constructors

Like !include will now give an error in the frontend.

* lint

homeassistant/components/lovelace/__init__.py

@@ -2,9 +2,10 @@
 import logging
 import uuid
 import os
-from os import O_WRONLY, O_CREAT, O_TRUNC
+from os import O_CREAT, O_TRUNC, O_WRONLY
 from collections import OrderedDict
-from typing import Union, List, Dict
+from typing import Dict, List, Union
+
 import voluptuous as vol
 
 from homeassistant.components import websocket_api
@@ -14,21 +15,45 @@ _LOGGER = logging.getLogger(__name__)
 DOMAIN = 'lovelace'
 REQUIREMENTS = ['ruamel.yaml==0.15.72']
 
+LOVELACE_CONFIG_FILE = 'ui-lovelace.yaml'
+JSON_TYPE = Union[List, Dict, str]  # pylint: disable=invalid-name
+
 OLD_WS_TYPE_GET_LOVELACE_UI = 'frontend/lovelace_config'
 WS_TYPE_GET_LOVELACE_UI = 'lovelace/config'
+WS_TYPE_GET_CARD = 'lovelace/config/card/get'
+WS_TYPE_SET_CARD = 'lovelace/config/card/set'
 
 SCHEMA_GET_LOVELACE_UI = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
     vol.Required('type'): vol.Any(WS_TYPE_GET_LOVELACE_UI,
                                   OLD_WS_TYPE_GET_LOVELACE_UI),
 })
 
-JSON_TYPE = Union[List, Dict, str]  # pylint: disable=invalid-name
+SCHEMA_GET_CARD = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_GET_CARD,
+    vol.Required('card_id'): str,
+    vol.Optional('format', default='yaml'): str,
+})
+
+SCHEMA_SET_CARD = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_SET_CARD,
+    vol.Required('card_id'): str,
+    vol.Required('card_config'): vol.Any(str, Dict),
+    vol.Optional('format', default='yaml'): str,
+})
 
 
 class WriteError(HomeAssistantError):
     """Error writing the data."""
 
 
+class CardNotFoundError(HomeAssistantError):
+    """Card not found in data."""
+
+
+class UnsupportedYamlError(HomeAssistantError):
+    """Unsupported YAML."""
+
+
 def save_yaml(fname: str, data: JSON_TYPE):
     """Save a YAML file."""
     from ruamel.yaml import YAML
@@ -45,7 +70,7 @@ def save_yaml(fname: str, data: JSON_TYPE):
         _LOGGER.error(str(exc))
         raise HomeAssistantError(exc)
     except OSError as exc:
-        _LOGGER.exception('Saving YAML file failed: %s', fname)
+        _LOGGER.exception('Saving YAML file %s failed: %s', fname, exc)
         raise WriteError(exc)
     finally:
         if os.path.exists(tmp_fname):
@@ -57,18 +82,29 @@ def save_yaml(fname: str, data: JSON_TYPE):
                 _LOGGER.error("YAML replacement cleanup failed: %s", exc)
 
 
+def _yaml_unsupported(loader, node):
+    raise UnsupportedYamlError(
+        'Unsupported YAML, you can not use {} in ui-lovelace.yaml'
+        .format(node.tag))
+
+
 def load_yaml(fname: str) -> JSON_TYPE:
     """Load a YAML file."""
     from ruamel.yaml import YAML
+    from ruamel.yaml.constructor import RoundTripConstructor
     from ruamel.yaml.error import YAMLError
+
+    RoundTripConstructor.add_constructor(None, _yaml_unsupported)
+
     yaml = YAML(typ='rt')
+
     try:
         with open(fname, encoding='utf-8') as conf_file:
             # If configuration file is empty YAML returns None
             # We convert that to an empty dict
             return yaml.load(conf_file) or OrderedDict()
     except YAMLError as exc:
-        _LOGGER.error("YAML error: %s", exc)
+        _LOGGER.error("YAML error in %s: %s", fname, exc)
         raise HomeAssistantError(exc)
     except UnicodeDecodeError as exc:
         _LOGGER.error("Unable to read file %s: %s", fname, exc)
@@ -76,21 +112,86 @@ def load_yaml(fname: str) -> JSON_TYPE:
 
 
 def load_config(fname: str) -> JSON_TYPE:
-    """Load a YAML file and adds id to card if not present."""
+    """Load a YAML file and adds id to views and cards if not present."""
     config = load_yaml(fname)
-    # Check if all cards have an ID or else add one
+    # Check if all views and cards have an id or else add one
     updated = False
+    index = 0
     for view in config.get('views', []):
+        if 'id' not in view:
+            updated = True
+            view.insert(0, 'id', index,
+                        comment="Automatically created id")
         for card in view.get('cards', []):
             if 'id' not in card:
                 updated = True
-                card['id'] = uuid.uuid4().hex
+                card.insert(0, 'id', uuid.uuid4().hex,
-                card.move_to_end('id', last=False)
+                            comment="Automatically created id")
+        index += 1
     if updated:
         save_yaml(fname, config)
     return config
 
 
+def object_to_yaml(data: JSON_TYPE) -> str:
+    """Create yaml string from object."""
+    from ruamel.yaml import YAML
+    from ruamel.yaml.error import YAMLError
+    from ruamel.yaml.compat import StringIO
+    yaml = YAML(typ='rt')
+    yaml.indent(sequence=4, offset=2)
+    stream = StringIO()
+    try:
+        yaml.dump(data, stream)
+        return stream.getvalue()
+    except YAMLError as exc:
+        _LOGGER.error("YAML error: %s", exc)
+        raise HomeAssistantError(exc)
+
+
+def yaml_to_object(data: str) -> JSON_TYPE:
+    """Create object from yaml string."""
+    from ruamel.yaml import YAML
+    from ruamel.yaml.error import YAMLError
+    yaml = YAML(typ='rt')
+    try:
+        return yaml.load(data)
+    except YAMLError as exc:
+        _LOGGER.error("YAML error: %s", exc)
+        raise HomeAssistantError(exc)
+
+
+def get_card(fname: str, card_id: str, data_format: str) -> JSON_TYPE:
+    """Load a specific card config for id."""
+    config = load_yaml(fname)
+    for view in config.get('views', []):
+        for card in view.get('cards', []):
+            if card.get('id') == card_id:
+                if data_format == 'yaml':
+                    return object_to_yaml(card)
+                return card
+
+    raise CardNotFoundError(
+        "Card with ID: {} was not found in {}.".format(card_id, fname))
+
+
+def set_card(fname: str, card_id: str, card_config: str, data_format: str)\
+        -> bool:
+    """Save a specific card config for id."""
+    config = load_yaml(fname)
+    for view in config.get('views', []):
+        for card in view.get('cards', []):
+            if card.get('id') == card_id:
+                if data_format == 'yaml':
+                    card_config = yaml_to_object(card_config)
+                card.update(card_config)
+                save_yaml(fname, config)
+                return True
+
+    raise CardNotFoundError(
+        "Card with ID: {} was not found in {}.".format(card_id, fname))
+
+
 async def async_setup(hass, config):
     """Set up the Lovelace commands."""
     # Backwards compat. Added in 0.80. Remove after 0.85
@@ -102,6 +203,14 @@ async def async_setup(hass, config):
         WS_TYPE_GET_LOVELACE_UI, websocket_lovelace_config,
         SCHEMA_GET_LOVELACE_UI)
 
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_GET_CARD, websocket_lovelace_get_card,
+        SCHEMA_GET_CARD)
+
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_SET_CARD, websocket_lovelace_set_card,
+        SCHEMA_SET_CARD)
+
     return True
 
 
@@ -111,13 +220,15 @@ async def websocket_lovelace_config(hass, connection, msg):
     error = None
     try:
         config = await hass.async_add_executor_job(
-            load_config, hass.config.path('ui-lovelace.yaml'))
+            load_config, hass.config.path(LOVELACE_CONFIG_FILE))
         message = websocket_api.result_message(
             msg['id'], config
         )
     except FileNotFoundError:
         error = ('file_not_found',
                  'Could not find ui-lovelace.yaml in your config dir.')
+    except UnsupportedYamlError as err:
+        error = 'unsupported_error', str(err)
     except HomeAssistantError as err:
         error = 'load_error', str(err)
 
@@ -125,3 +236,59 @@ async def websocket_lovelace_config(hass, connection, msg):
         message = websocket_api.error_message(msg['id'], *error)
 
     connection.send_message(message)
+
+
+@websocket_api.async_response
+async def websocket_lovelace_get_card(hass, connection, msg):
+    """Send lovelace card config over websocket config."""
+    error = None
+    try:
+        card = await hass.async_add_executor_job(
+            get_card, hass.config.path(LOVELACE_CONFIG_FILE), msg['card_id'],
+            msg.get('format', 'yaml'))
+        message = websocket_api.result_message(
+            msg['id'], card
+        )
+    except FileNotFoundError:
+        error = ('file_not_found',
+                 'Could not find ui-lovelace.yaml in your config dir.')
+    except UnsupportedYamlError as err:
+        error = 'unsupported_error', str(err)
+    except CardNotFoundError:
+        error = ('card_not_found',
+                 'Could not find card in ui-lovelace.yaml.')
+    except HomeAssistantError as err:
+        error = 'load_error', str(err)
+
+    if error is not None:
+        message = websocket_api.error_message(msg['id'], *error)
+
+    connection.send_message(message)
+
+
+@websocket_api.async_response
+async def websocket_lovelace_set_card(hass, connection, msg):
+    """Receive lovelace card config over websocket and save."""
+    error = None
+    try:
+        result = await hass.async_add_executor_job(
+            set_card, hass.config.path(LOVELACE_CONFIG_FILE),
+            msg['card_id'], msg['card_config'], msg.get('format', 'yaml'))
+        message = websocket_api.result_message(
+            msg['id'], result
+        )
+    except FileNotFoundError:
+        error = ('file_not_found',
+                 'Could not find ui-lovelace.yaml in your config dir.')
+    except UnsupportedYamlError as err:
+        error = 'unsupported_error', str(err)
+    except CardNotFoundError:
+        error = ('card_not_found',
+                 'Could not find card in ui-lovelace.yaml.')
+    except HomeAssistantError as err:
+        error = 'save_error', str(err)
+
+    if error is not None:
+        message = websocket_api.error_message(msg['id'], *error)
+
+    connection.send_message(message)

tests/components/lovelace/test_init.py

@@ -9,7 +9,8 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.setup import async_setup_component
 from homeassistant.components.websocket_api.const import TYPE_RESULT
 from homeassistant.components.lovelace import (load_yaml,
-                                               save_yaml, load_config)
+                                               save_yaml, load_config,
+                                               UnsupportedYamlError)
 
 TEST_YAML_A = """\
 title: My Awesome Home
@@ -55,6 +56,8 @@ views:
     # Title of the view. Will be used as the tooltip for tab icon
     title: Second view
     cards:
+      - id: test
+        type: entities
         # Entities card will take a list of entities and show their state.
       - type: entities
         # Title of the entities card
@@ -79,6 +82,7 @@ TEST_YAML_B = """\
 title: Home
 views:
   - title: Dashboard
+    id: dashboard
     icon: mdi:home
     cards:
       - id: testid
@@ -102,6 +106,15 @@ views:
           type: vertical-stack
 """
 
+# Test unsupported YAML
+TEST_UNSUP_YAML = """\
+title: Home
+views:
+  - title: Dashboard
+    icon: mdi:home
+    cards: !include cards.yaml
+"""
+
 
 class TestYAML(unittest.TestCase):
     """Test lovelace.yaml save and load."""
@@ -147,9 +160,11 @@ class TestYAML(unittest.TestCase):
         """Test if id is added."""
         fname = self._path_for("test6")
         with patch('homeassistant.components.lovelace.load_yaml',
-                   return_value=self.yaml.load(TEST_YAML_A)):
+                   return_value=self.yaml.load(TEST_YAML_A)), \
+                patch('homeassistant.components.lovelace.save_yaml'):
             data = load_config(fname)
         assert 'id' in data['views'][0]['cards'][0]
+        assert 'id' in data['views'][1]
 
     def test_id_not_changed(self):
         """Test if id is not changed if already exists."""
@@ -256,7 +271,7 @@ async def test_lovelace_ui_not_found(hass, hass_ws_client):
 
 
 async def test_lovelace_ui_load_err(hass, hass_ws_client):
-    """Test lovelace_ui command cannot find file."""
+    """Test lovelace_ui command load error."""
     await async_setup_component(hass, 'lovelace')
     client = await hass_ws_client(hass)
 
@@ -272,3 +287,156 @@ async def test_lovelace_ui_load_err(hass, hass_ws_client):
     assert msg['type'] == TYPE_RESULT
     assert msg['success'] is False
     assert msg['error']['code'] == 'load_error'
+
+
+async def test_lovelace_ui_load_json_err(hass, hass_ws_client):
+    """Test lovelace_ui command load error."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+
+    with patch('homeassistant.components.lovelace.load_config',
+               side_effect=UnsupportedYamlError):
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config',
+        })
+        msg = await client.receive_json()
+
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success'] is False
+    assert msg['error']['code'] == 'unsupported_error'
+
+
+async def test_lovelace_get_card(hass, hass_ws_client):
+    """Test get_card command."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+    yaml = YAML(typ='rt')
+
+    with patch('homeassistant.components.lovelace.load_yaml',
+               return_value=yaml.load(TEST_YAML_A)):
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config/card/get',
+            'card_id': 'test',
+        })
+        msg = await client.receive_json()
+
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success']
+    assert msg['result'] == 'id: test\ntype: entities\n'
+
+
+async def test_lovelace_get_card_not_found(hass, hass_ws_client):
+    """Test get_card command cannot find card."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+    yaml = YAML(typ='rt')
+
+    with patch('homeassistant.components.lovelace.load_yaml',
+               return_value=yaml.load(TEST_YAML_A)):
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config/card/get',
+            'card_id': 'not_found',
+        })
+        msg = await client.receive_json()
+
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success'] is False
+    assert msg['error']['code'] == 'card_not_found'
+
+
+async def test_lovelace_get_card_bad_yaml(hass, hass_ws_client):
+    """Test get_card command bad yaml."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+
+    with patch('homeassistant.components.lovelace.load_yaml',
+               side_effect=HomeAssistantError):
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config/card/get',
+            'card_id': 'testid',
+        })
+        msg = await client.receive_json()
+
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success'] is False
+    assert msg['error']['code'] == 'load_error'
+
+
+async def test_lovelace_set_card(hass, hass_ws_client):
+    """Test set_card command."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+    yaml = YAML(typ='rt')
+
+    with patch('homeassistant.components.lovelace.load_yaml',
+               return_value=yaml.load(TEST_YAML_A)), \
+        patch('homeassistant.components.lovelace.save_yaml') \
+            as save_yaml_mock:
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config/card/set',
+            'card_id': 'test',
+            'card_config': 'id: test\ntype: glance\n',
+        })
+        msg = await client.receive_json()
+
+    result = save_yaml_mock.call_args_list[0][0][1]
+    assert result.mlget(['views', 1, 'cards', 0, 'type'],
+                        list_ok=True) == 'glance'
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success']
+
+
+async def test_lovelace_set_card_not_found(hass, hass_ws_client):
+    """Test set_card command cannot find card."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+    yaml = YAML(typ='rt')
+
+    with patch('homeassistant.components.lovelace.load_yaml',
+               return_value=yaml.load(TEST_YAML_A)):
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config/card/set',
+            'card_id': 'not_found',
+            'card_config': 'id: test\ntype: glance\n',
+        })
+        msg = await client.receive_json()
+
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success'] is False
+    assert msg['error']['code'] == 'card_not_found'
+
+
+async def test_lovelace_set_card_bad_yaml(hass, hass_ws_client):
+    """Test set_card command bad yaml."""
+    await async_setup_component(hass, 'lovelace')
+    client = await hass_ws_client(hass)
+    yaml = YAML(typ='rt')
+
+    with patch('homeassistant.components.lovelace.load_yaml',
+               return_value=yaml.load(TEST_YAML_A)), \
+        patch('homeassistant.components.lovelace.yaml_to_object',
+              side_effect=HomeAssistantError):
+        await client.send_json({
+            'id': 5,
+            'type': 'lovelace/config/card/set',
+            'card_id': 'test',
+            'card_config': 'id: test\ntype: glance\n',
+        })
+        msg = await client.receive_json()
+
+    assert msg['id'] == 5
+    assert msg['type'] == TYPE_RESULT
+    assert msg['success'] is False
+    assert msg['error']['code'] == 'save_error'