From 00b8d57cd0cf64daa0ddf59fa6cf9da1af824d8f Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Tue, 12 Feb 2019 07:38:19 -0800 Subject: [PATCH] Add frontend storage (#20880) * Add frontend storage * Update storage.py --- homeassistant/components/frontend/__init__.py | 3 + homeassistant/components/frontend/storage.py | 79 ++++++++ tests/components/frontend/test_storage.py | 186 ++++++++++++++++++ 3 files changed, 268 insertions(+) create mode 100644 homeassistant/components/frontend/storage.py create mode 100644 tests/components/frontend/test_storage.py diff --git a/homeassistant/components/frontend/__init__.py b/homeassistant/components/frontend/__init__.py index 46652b4d7b0..1efda81dfe0 100644 --- a/homeassistant/components/frontend/__init__.py +++ b/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( diff --git a/homeassistant/components/frontend/storage.py b/homeassistant/components/frontend/storage.py new file mode 100644 index 00000000000..f01abc79e8e --- /dev/null +++ b/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 + } + )) diff --git a/tests/components/frontend/test_storage.py b/tests/components/frontend/test_storage.py new file mode 100644 index 00000000000..97b132cfd13 --- /dev/null +++ b/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'