Add frontend storage (#20880)

* Add frontend storage

* Update storage.py

homeassistant/components/frontend/__init__.py

@@ -24,6 +24,8 @@ from homeassistant.core import callback
 from homeassistant.helpers.translation import async_get_translations
 from homeassistant.loader import bind_hass
 
+from .storage import async_setup_frontend_storage
+
 REQUIREMENTS = ['home-assistant-frontend==20190203.0']
 
 DOMAIN = 'frontend'
@@ -195,6 +197,7 @@ def add_manifest_json_key(key, val):
 
 async def async_setup(hass, config):
     """Set up the serving of the frontend."""
+    await async_setup_frontend_storage(hass)
     hass.components.websocket_api.async_register_command(
         WS_TYPE_GET_PANELS, websocket_get_panels, SCHEMA_GET_PANELS)
     hass.components.websocket_api.async_register_command(

homeassistant/components/frontend/storage.py

@@ -0,0 +1,79 @@
+"""API for persistent storage for the frontend."""
+from functools import wraps
+import voluptuous as vol
+
+from homeassistant.components import websocket_api
+
+DATA_STORAGE = 'frontend_storage'
+STORAGE_VERSION_USER_DATA = 1
+STORAGE_KEY_USER_DATA = 'frontend.user_data_{}'
+
+
+async def async_setup_frontend_storage(hass):
+    """Set up frontend storage."""
+    hass.data[DATA_STORAGE] = {}
+    hass.components.websocket_api.async_register_command(
+        websocket_set_user_data
+    )
+    hass.components.websocket_api.async_register_command(
+        websocket_get_user_data
+    )
+
+
+def with_store(orig_func):
+    """Decorate function to provide data."""
+    @wraps(orig_func)
+    async def with_store_func(hass, connection, msg):
+        """Provide user specific data and store to function."""
+        store = hass.helpers.storage.Store(
+            STORAGE_VERSION_USER_DATA,
+            STORAGE_KEY_USER_DATA.format(connection.user.id)
+        )
+        data = hass.data[DATA_STORAGE]
+        user_id = connection.user.id
+        if user_id not in data:
+            data[user_id] = await store.async_load() or {}
+
+        await orig_func(
+            hass, connection, msg,
+            store,
+            data[user_id],
+        )
+    return with_store_func
+
+
+@websocket_api.websocket_command({
+    vol.Required('type'): 'frontend/set_user_data',
+    vol.Required('key'): str,
+    vol.Required('value'): vol.Any(bool, str, int, float, dict, list, None),
+})
+@websocket_api.async_response
+@with_store
+async def websocket_set_user_data(hass, connection, msg, store, data):
+    """Handle set global data command.
+
+    Async friendly.
+    """
+    data[msg['key']] = msg['value']
+    await store.async_save(data)
+    connection.send_message(websocket_api.result_message(
+        msg['id'],
+    ))
+
+
+@websocket_api.websocket_command({
+    vol.Required('type'): 'frontend/get_user_data',
+    vol.Optional('key'): str,
+})
+@websocket_api.async_response
+@with_store
+async def websocket_get_user_data(hass, connection, msg, store, data):
+    """Handle get global data command.
+
+    Async friendly.
+    """
+    connection.send_message(websocket_api.result_message(
+        msg['id'], {
+            'value': data.get(msg['key']) if 'key' in msg else data
+        }
+    ))

tests/components/frontend/test_storage.py

@@ -0,0 +1,186 @@
+"""The tests for frontend storage."""
+import pytest
+
+from homeassistant.setup import async_setup_component
+from homeassistant.components.frontend import storage
+
+
+@pytest.fixture(autouse=True)
+def setup_frontend(hass):
+    """Fixture to setup the frontend."""
+    hass.loop.run_until_complete(async_setup_component(hass, 'frontend', {}))
+
+
+async def test_get_user_data_empty(hass, hass_ws_client, hass_storage):
+    """Test get_user_data command."""
+    client = await hass_ws_client(hass)
+
+    await client.send_json({
+        'id': 5,
+        'type': 'frontend/get_user_data',
+        'key': 'non-existing-key',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'] is None
+
+
+async def test_get_user_data(hass, hass_ws_client, hass_admin_user,
+                             hass_storage):
+    """Test get_user_data command."""
+    storage_key = storage.STORAGE_KEY_USER_DATA.format(hass_admin_user.id)
+    hass_storage[storage_key] = {
+        'key': storage_key,
+        'version': 1,
+        'data': {
+            'test-key': 'test-value',
+            'test-complex': [{'foo': 'bar'}]
+        }
+    }
+
+    client = await hass_ws_client(hass)
+
+    # Get a simple string key
+
+    await client.send_json({
+        'id': 6,
+        'type': 'frontend/get_user_data',
+        'key': 'test-key',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'] == 'test-value'
+
+    # Get a more complex key
+
+    await client.send_json({
+        'id': 7,
+        'type': 'frontend/get_user_data',
+        'key': 'test-complex',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'][0]['foo'] == 'bar'
+
+    # Get all data (no key)
+
+    await client.send_json({
+        'id': 8,
+        'type': 'frontend/get_user_data',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value']['test-key'] == 'test-value'
+    assert res['result']['value']['test-complex'][0]['foo'] == 'bar'
+
+
+async def test_set_user_data_empty(hass, hass_ws_client, hass_storage):
+    """Test set_user_data command."""
+    client = await hass_ws_client(hass)
+
+    # test creating
+
+    await client.send_json({
+        'id': 6,
+        'type': 'frontend/get_user_data',
+        'key': 'test-key',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'] is None
+
+    await client.send_json({
+        'id': 7,
+        'type': 'frontend/set_user_data',
+        'key': 'test-key',
+        'value': 'test-value'
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+
+    await client.send_json({
+        'id': 8,
+        'type': 'frontend/get_user_data',
+        'key': 'test-key',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'] == 'test-value'
+
+
+async def test_set_user_data(hass, hass_ws_client, hass_storage,
+                             hass_admin_user):
+    """Test set_user_data command with initial data."""
+    storage_key = storage.STORAGE_KEY_USER_DATA.format(hass_admin_user.id)
+    hass_storage[storage_key] = {
+        'version': 1,
+        'data': {
+            'test-key': 'test-value',
+            'test-complex': 'string',
+        }
+    }
+
+    client = await hass_ws_client(hass)
+
+    # test creating
+
+    await client.send_json({
+        'id': 5,
+        'type': 'frontend/set_user_data',
+        'key': 'test-non-existent-key',
+        'value': 'test-value-new'
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+
+    await client.send_json({
+        'id': 6,
+        'type': 'frontend/get_user_data',
+        'key': 'test-non-existent-key',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'] == 'test-value-new'
+
+    # test updating with complex data
+
+    await client.send_json({
+        'id': 7,
+        'type': 'frontend/set_user_data',
+        'key': 'test-complex',
+        'value': [{'foo': 'bar'}]
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+
+    await client.send_json({
+        'id': 8,
+        'type': 'frontend/get_user_data',
+        'key': 'test-complex',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'][0]['foo'] == 'bar'
+
+    # ensure other existing key was not modified
+
+    await client.send_json({
+        'id': 9,
+        'type': 'frontend/get_user_data',
+        'key': 'test-key',
+    })
+
+    res = await client.receive_json()
+    assert res['success'], res
+    assert res['result']['value'] == 'test-value'