Area registry (#20435)

* First draft of area registry

* Refactor based on input

* Add tests for areas
Add tests for updating device

* Updating a device shouldn't require area

* Fix Martins comment

* Require admin

* Save after deleting

* Rename read to list_areas
Fix device entry_dict
Remove area id from device when deleting area

* Fix tests

homeassistant/components/config/area_registry.py

@@ -0,0 +1,126 @@
+"""HTTP views to interact with the area registry."""
+import voluptuous as vol
+
+from homeassistant.components import websocket_api
+from homeassistant.components.websocket_api.decorators import (
+    async_response, require_admin)
+from homeassistant.core import callback
+from homeassistant.helpers.area_registry import async_get_registry
+
+
+DEPENDENCIES = ['websocket_api']
+
+WS_TYPE_LIST = 'config/area_registry/list'
+SCHEMA_WS_LIST = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_LIST,
+})
+
+WS_TYPE_CREATE = 'config/area_registry/create'
+SCHEMA_WS_CREATE = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_CREATE,
+    vol.Required('name'): str,
+})
+
+WS_TYPE_DELETE = 'config/area_registry/delete'
+SCHEMA_WS_DELETE = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_DELETE,
+    vol.Required('area_id'): str,
+})
+
+WS_TYPE_UPDATE = 'config/area_registry/update'
+SCHEMA_WS_UPDATE = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_UPDATE,
+    vol.Required('area_id'): str,
+    vol.Required('name'): str,
+})
+
+
+async def async_setup(hass):
+    """Enable the Area Registry views."""
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_LIST, websocket_list_areas, SCHEMA_WS_LIST
+    )
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_CREATE, websocket_create_area, SCHEMA_WS_CREATE
+    )
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_DELETE, websocket_delete_area, SCHEMA_WS_DELETE
+    )
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_UPDATE, websocket_update_area, SCHEMA_WS_UPDATE
+    )
+    return True
+
+
+@async_response
+async def websocket_list_areas(hass, connection, msg):
+    """Handle list areas command."""
+    registry = await async_get_registry(hass)
+    connection.send_message(websocket_api.result_message(
+        msg['id'], [{
+            'name': entry.name,
+            'area_id': entry.id,
+        } for entry in registry.async_list_areas()]
+    ))
+
+
+@require_admin
+@async_response
+async def websocket_create_area(hass, connection, msg):
+    """Create area command."""
+    registry = await async_get_registry(hass)
+    try:
+        entry = registry.async_create(msg['name'])
+    except ValueError as err:
+        connection.send_message(websocket_api.error_message(
+            msg['id'], 'invalid_info', str(err)
+        ))
+    else:
+        connection.send_message(websocket_api.result_message(
+            msg['id'], _entry_dict(entry)
+        ))
+
+
+@require_admin
+@async_response
+async def websocket_delete_area(hass, connection, msg):
+    """Delete area command."""
+    registry = await async_get_registry(hass)
+
+    try:
+        await registry.async_delete(msg['area_id'])
+    except KeyError:
+        connection.send_message(websocket_api.error_message(
+            msg['id'], 'invalid_info', "Area ID doesn't exist"
+        ))
+    else:
+        connection.send_message(websocket_api.result_message(
+            msg['id'], 'success'
+        ))
+
+
+@require_admin
+@async_response
+async def websocket_update_area(hass, connection, msg):
+    """Handle update area websocket command."""
+    registry = await async_get_registry(hass)
+
+    try:
+        entry = registry.async_update(msg['area_id'], msg['name'])
+    except ValueError as err:
+        connection.send_message(websocket_api.error_message(
+            msg['id'], 'invalid_info', str(err)
+        ))
+    else:
+        connection.send_message(websocket_api.result_message(
+            msg['id'], _entry_dict(entry)
+        ))
+
+
+@callback
+def _entry_dict(entry):
+    """Convert entry to API format."""
+    return {
+        'area_id': entry.id,
+        'name': entry.name
+    }

homeassistant/components/config/device_registry.py

@@ -1,8 +1,11 @@
 """HTTP views to interact with the device registry."""
 import voluptuous as vol
 
-from homeassistant.helpers.device_registry import async_get_registry
 from homeassistant.components import websocket_api
+from homeassistant.components.websocket_api.decorators import (
+    async_response, require_admin)
+from homeassistant.core import callback
+from homeassistant.helpers.device_registry import async_get_registry
 
 DEPENDENCIES = ['websocket_api']
 
@@ -11,22 +14,53 @@ SCHEMA_WS_LIST = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
     vol.Required('type'): WS_TYPE_LIST,
 })
 
+WS_TYPE_UPDATE = 'config/device_registry/update'
+SCHEMA_WS_UPDATE = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_UPDATE,
+    vol.Required('device_id'): str,
+    vol.Optional('area_id'): str,
+})
+
 
 async def async_setup(hass):
-    """Enable the Entity Registry views."""
+    """Enable the Device Registry views."""
     hass.components.websocket_api.async_register_command(
         WS_TYPE_LIST, websocket_list_devices,
         SCHEMA_WS_LIST
     )
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_UPDATE, websocket_update_device, SCHEMA_WS_UPDATE
+    )
     return True
 
 
-@websocket_api.async_response
+@async_response
 async def websocket_list_devices(hass, connection, msg):
     """Handle list devices command."""
     registry = await async_get_registry(hass)
     connection.send_message(websocket_api.result_message(
-        msg['id'], [{
+        msg['id'], [_entry_dict(entry) for entry in registry.devices.values()]
+    ))
+
+
+@require_admin
+@async_response
+async def websocket_update_device(hass, connection, msg):
+    """Handle update area websocket command."""
+    registry = await async_get_registry(hass)
+
+    entry = registry.async_update_device(
+        msg['device_id'], area_id=msg['area_id'])
+
+    connection.send_message(websocket_api.result_message(
+        msg['id'], _entry_dict(entry)
+    ))
+
+
+@callback
+def _entry_dict(entry):
+    """Convert entry to API format."""
+    return {
         'config_entries': list(entry.config_entries),
         'connections': list(entry.connections),
         'manufacturer': entry.manufacturer,
@@ -35,5 +69,5 @@ async def websocket_list_devices(hass, connection, msg):
         'sw_version': entry.sw_version,
         'id': entry.id,
         'hub_device_id': entry.hub_device_id,
-        } for entry in registry.devices.values()]
+        'area_id': entry.area_id,
-    ))
+    }

homeassistant/helpers/area_registry.py

@@ -0,0 +1,139 @@
+"""Provide a way to connect devices to one physical location."""
+import logging
+import uuid
+from collections import OrderedDict
+from typing import List, Optional
+
+import attr
+
+from homeassistant.core import callback
+from homeassistant.loader import bind_hass
+
+_LOGGER = logging.getLogger(__name__)
+
+DATA_REGISTRY = 'area_registry'
+
+STORAGE_KEY = 'core.area_registry'
+STORAGE_VERSION = 1
+SAVE_DELAY = 10
+
+
+@attr.s(slots=True, frozen=True)
+class AreaEntry:
+    """Area Registry Entry."""
+
+    name = attr.ib(type=str, default=None)
+    id = attr.ib(type=str, default=attr.Factory(lambda: uuid.uuid4().hex))
+
+
+class AreaRegistry:
+    """Class to hold a registry of areas."""
+
+    def __init__(self, hass) -> None:
+        """Initialize the area registry."""
+        self.hass = hass
+        self.areas = None
+        self._store = hass.helpers.storage.Store(STORAGE_VERSION, STORAGE_KEY)
+
+    @callback
+    def async_list_areas(self) -> List[AreaEntry]:
+        """Get all areas."""
+        return self.areas.values()
+
+    @callback
+    def async_create(self, name: str) -> AreaEntry:
+        """Create a new area."""
+        if self._async_is_registered(name):
+            raise ValueError('Name is already in use')
+
+        area = AreaEntry()
+        self.areas[area.id] = area
+
+        return self.async_update(area.id, name=name)
+
+    async def async_delete(self, area_id: str) -> None:
+        """Delete area."""
+        device_registry = await \
+            self.hass.helpers.device_registry.async_get_registry()
+        device_registry.async_clear_area_id(area_id)
+
+        del self.areas[area_id]
+
+        self.async_schedule_save()
+
+    @callback
+    def async_update(self, area_id: str, name: str) -> AreaEntry:
+        """Update name of area."""
+        old = self.areas[area_id]
+
+        changes = {}
+
+        if name == old.name:
+            return old
+
+        if self._async_is_registered(name):
+            raise ValueError('Name is already in use')
+        else:
+            changes['name'] = name
+
+        new = self.areas[area_id] = attr.evolve(old, **changes)
+        self.async_schedule_save()
+        return new
+
+    @callback
+    def _async_is_registered(self, name) -> Optional[AreaEntry]:
+        """Check if a name is currently registered."""
+        for area in self.areas.values():
+            if name == area.name:
+                return area
+        return False
+
+    async def async_load(self) -> None:
+        """Load the area registry."""
+        data = await self._store.async_load()
+
+        areas = OrderedDict()
+
+        if data is not None:
+            for area in data['areas']:
+                areas[area['id']] = AreaEntry(
+                    name=area['name'],
+                    id=area['id']
+                )
+
+        self.areas = areas
+
+    @callback
+    def async_schedule_save(self) -> None:
+        """Schedule saving the area registry."""
+        self._store.async_delay_save(self._data_to_save, SAVE_DELAY)
+
+    @callback
+    def _data_to_save(self) -> dict:
+        """Return data of area registry to store in a file."""
+        data = {}
+
+        data['areas'] = [
+            {
+                'name': entry.name,
+                'id': entry.id,
+            } for entry in self.areas.values()
+        ]
+
+        return data
+
+
+@bind_hass
+async def async_get_registry(hass) -> AreaRegistry:
+    """Return area registry instance."""
+    task = hass.data.get(DATA_REGISTRY)
+
+    if task is None:
+        async def _load_reg():
+            registry = AreaRegistry(hass)
+            await registry.async_load()
+            return registry
+
+        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+
+    return await task

homeassistant/helpers/device_registry.py

@@ -36,6 +36,7 @@ class DeviceEntry:
     name = attr.ib(type=str, default=None)
     sw_version = attr.ib(type=str, default=None)
     hub_device_id = attr.ib(type=str, default=None)
+    area_id = attr.ib(type=str, default=None)
     id = attr.ib(type=str, default=attr.Factory(lambda: uuid.uuid4().hex))
 
 
@@ -119,9 +120,14 @@ class DeviceRegistry:
             manufacturer=manufacturer,
             model=model,
             name=name,
-            sw_version=sw_version,
+            sw_version=sw_version
         )
 
+    @callback
+    def async_update_device(self, device_id, *, area_id=_UNDEF):
+        """Update properties of a device."""
+        return self._async_update_device(device_id, area_id=area_id)
+
     @callback
     def _async_update_device(self, device_id, *, add_config_entry_id=_UNDEF,
                              remove_config_entry_id=_UNDEF,
@@ -131,7 +137,8 @@ class DeviceRegistry:
                              model=_UNDEF,
                              name=_UNDEF,
                              sw_version=_UNDEF,
-                             hub_device_id=_UNDEF):
+                             hub_device_id=_UNDEF,
+                             area_id=_UNDEF):
         """Update device attributes."""
         old = self.devices[device_id]
 
@@ -169,6 +176,9 @@ class DeviceRegistry:
             if value is not _UNDEF and value != getattr(old, attr_name):
                 changes[attr_name] = value
 
+        if (area_id is not _UNDEF and area_id != old.area_id):
+            changes['area_id'] = area_id
+
         if not changes:
             return old
 
@@ -197,6 +207,8 @@ class DeviceRegistry:
                     id=device['id'],
                     # Introduced in 0.79
                     hub_device_id=device.get('hub_device_id'),
+                    # Introduced in 0.87
+                    area_id=device.get('area_id')
                 )
 
         self.devices = devices
@@ -222,6 +234,7 @@ class DeviceRegistry:
                 'sw_version': entry.sw_version,
                 'id': entry.id,
                 'hub_device_id': entry.hub_device_id,
+                'area_id': entry.area_id
             } for entry in self.devices.values()
         ]
 
@@ -235,6 +248,13 @@ class DeviceRegistry:
                 self._async_update_device(
                     dev_id, remove_config_entry_id=config_entry_id)
 
+    @callback
+    def async_clear_area_id(self, area_id: str) -> None:
+        """Clear area id from registry entries."""
+        for dev_id, device in self.devices.items():
+            if area_id == device.area_id:
+                self._async_update_device(dev_id, area_id=None)
+
 
 @bind_hass
 async def async_get_registry(hass) -> DeviceRegistry:

tests/common.py

@@ -1,35 +1,37 @@
 """Test the helper method for writing tests."""
 import asyncio
-from collections import OrderedDict
-from datetime import timedelta
 import functools as ft
 import json
+import logging
 import os
 import sys
-from unittest.mock import patch, MagicMock, Mock
-from io import StringIO
-import logging
 import threading
-from contextlib import contextmanager
 
-from homeassistant import auth, core as ha, config_entries
+from collections import OrderedDict
+from contextlib import contextmanager
+from datetime import timedelta
+from io import StringIO
+from unittest.mock import MagicMock, Mock, patch
+
+import homeassistant.util.dt as date_util
+import homeassistant.util.yaml as yaml
+
+from homeassistant import auth, config_entries, core as ha
 from homeassistant.auth import (
     models as auth_models, auth_store, providers as auth_providers,
     permissions as auth_permissions)
 from homeassistant.auth.permissions import system_policies
-from homeassistant.setup import setup_component, async_setup_component
-from homeassistant.config import async_process_component_config
-from homeassistant.helpers import (
-    intent, entity, restore_state, entity_registry,
-    entity_platform, storage, device_registry)
-from homeassistant.util.unit_system import METRIC_SYSTEM
-import homeassistant.util.dt as date_util
-import homeassistant.util.yaml as yaml
-from homeassistant.const import (
-    STATE_ON, STATE_OFF, DEVICE_DEFAULT_NAME, EVENT_TIME_CHANGED,
-    EVENT_STATE_CHANGED, EVENT_PLATFORM_DISCOVERED, ATTR_SERVICE,
-    ATTR_DISCOVERED, SERVER_PORT, EVENT_HOMEASSISTANT_CLOSE)
 from homeassistant.components import mqtt, recorder
+from homeassistant.config import async_process_component_config
+from homeassistant.const import (
+    ATTR_DISCOVERED, ATTR_SERVICE, DEVICE_DEFAULT_NAME,
+    EVENT_HOMEASSISTANT_CLOSE, EVENT_PLATFORM_DISCOVERED, EVENT_STATE_CHANGED,
+    EVENT_TIME_CHANGED, SERVER_PORT, STATE_ON, STATE_OFF)
+from homeassistant.helpers import (
+    area_registry, device_registry, entity, entity_platform, entity_registry,
+    intent, restore_state, storage)
+from homeassistant.setup import async_setup_component, setup_component
+from homeassistant.util.unit_system import METRIC_SYSTEM
 from homeassistant.util.async_ import (
     run_callback_threadsafe, run_coroutine_threadsafe)
 
@@ -333,6 +335,19 @@ def mock_registry(hass, mock_entries=None):
     return registry
 
 
+def mock_area_registry(hass, mock_entries=None):
+    """Mock the Area Registry."""
+    registry = area_registry.AreaRegistry(hass)
+    registry.areas = mock_entries or OrderedDict()
+
+    async def _get_reg():
+        return registry
+
+    hass.data[area_registry.DATA_REGISTRY] = \
+        hass.loop.create_task(_get_reg())
+    return registry
+
+
 def mock_device_registry(hass, mock_entries=None):
     """Mock the Device Registry."""
     registry = device_registry.DeviceRegistry(hass)

tests/components/config/test_area_registry.py

@@ -0,0 +1,155 @@
+"""Test area_registry API."""
+import pytest
+
+from homeassistant.components.config import area_registry
+from tests.common import mock_area_registry
+
+
+@pytest.fixture
+def client(hass, hass_ws_client):
+    """Fixture that can interact with the config manager API."""
+    hass.loop.run_until_complete(area_registry.async_setup(hass))
+    yield hass.loop.run_until_complete(hass_ws_client(hass))
+
+
+@pytest.fixture
+def registry(hass):
+    """Return an empty, loaded, registry."""
+    return mock_area_registry(hass)
+
+
+async def test_list_areas(hass, client, registry):
+    """Test list entries."""
+    registry.async_create('mock 1')
+    registry.async_create('mock 2')
+
+    await client.send_json({
+        'id': 1,
+        'type': 'config/area_registry/list',
+    })
+
+    msg = await client.receive_json()
+
+    assert len(msg['result']) == len(registry.areas)
+
+
+async def test_create_area(hass, client, registry):
+    """Test create entry."""
+    await client.send_json({
+        'id': 1,
+        'name': "mock",
+        'type': 'config/area_registry/create',
+    })
+
+    msg = await client.receive_json()
+
+    assert 'mock' in msg['result']['name']
+    assert len(registry.areas) == 1
+
+
+async def test_create_area_with_name_already_in_use(hass, client, registry):
+    """Test create entry that should fail."""
+    registry.async_create('mock')
+
+    await client.send_json({
+        'id': 1,
+        'name': "mock",
+        'type': 'config/area_registry/create',
+    })
+
+    msg = await client.receive_json()
+
+    assert not msg['success']
+    assert msg['error']['code'] == 'invalid_info'
+    assert msg['error']['message'] == "Name is already in use"
+    assert len(registry.areas) == 1
+
+
+async def test_delete_area(hass, client, registry):
+    """Test delete entry."""
+    area = registry.async_create('mock')
+
+    await client.send_json({
+        'id': 1,
+        'area_id': area.id,
+        'type': 'config/area_registry/delete',
+    })
+
+    msg = await client.receive_json()
+
+    assert msg['success']
+    assert not registry.areas
+
+
+async def test_delete_non_existing_area(hass, client, registry):
+    """Test delete entry that should fail."""
+    registry.async_create('mock')
+
+    await client.send_json({
+        'id': 1,
+        'area_id': '',
+        'type': 'config/area_registry/delete',
+    })
+
+    msg = await client.receive_json()
+
+    assert not msg['success']
+    assert msg['error']['code'] == 'invalid_info'
+    assert msg['error']['message'] == "Area ID doesn't exist"
+    assert len(registry.areas) == 1
+
+
+async def test_update_area(hass, client, registry):
+    """Test update entry."""
+    area = registry.async_create('mock 1')
+
+    await client.send_json({
+        'id': 1,
+        'area_id': area.id,
+        'name': "mock 2",
+        'type': 'config/area_registry/update',
+    })
+
+    msg = await client.receive_json()
+
+    assert msg['result']['area_id'] == area.id
+    assert msg['result']['name'] == 'mock 2'
+    assert len(registry.areas) == 1
+
+
+async def test_update_area_with_same_name(hass, client, registry):
+    """Test update entry."""
+    area = registry.async_create('mock 1')
+
+    await client.send_json({
+        'id': 1,
+        'area_id': area.id,
+        'name': "mock 1",
+        'type': 'config/area_registry/update',
+    })
+
+    msg = await client.receive_json()
+
+    assert msg['result']['area_id'] == area.id
+    assert msg['result']['name'] == 'mock 1'
+    assert len(registry.areas) == 1
+
+
+async def test_update_area_with_name_already_in_use(hass, client, registry):
+    """Test update entry."""
+    area = registry.async_create('mock 1')
+    registry.async_create('mock 2')
+
+    await client.send_json({
+        'id': 1,
+        'area_id': area.id,
+        'name': "mock 2",
+        'type': 'config/area_registry/update',
+    })
+
+    msg = await client.receive_json()
+
+    assert not msg['success']
+    assert msg['error']['code'] == 'invalid_info'
+    assert msg['error']['message'] == "Name is already in use"
+    assert len(registry.areas) == 2

tests/components/config/test_device_registry.py

@@ -1,4 +1,4 @@
-"""Test entity_registry API."""
+"""Test device_registry API."""
 import pytest
 
 from homeassistant.components.config import device_registry
@@ -48,6 +48,7 @@ async def test_list_devices(hass, client, registry):
             'name': None,
             'sw_version': None,
             'hub_device_id': None,
+            'area_id': None,
         },
         {
             'config_entries': ['1234'],
@@ -57,5 +58,30 @@ async def test_list_devices(hass, client, registry):
             'name': None,
             'sw_version': None,
             'hub_device_id': dev1,
+            'area_id': None,
         }
     ]
+
+
+async def test_update_device(hass, client, registry):
+    """Test update entry."""
+    device = registry.async_get_or_create(
+        config_entry_id='1234',
+        connections={('ethernet', '12:34:56:78:90:AB:CD:EF')},
+        identifiers={('bridgeid', '0123')},
+        manufacturer='manufacturer', model='model')
+
+    assert not device.area_id
+
+    await client.send_json({
+        'id': 1,
+        'device_id': device.id,
+        'area_id': '12345A',
+        'type': 'config/device_registry/update',
+    })
+
+    msg = await client.receive_json()
+
+    assert msg['result']['id'] == device.id
+    assert msg['result']['area_id'] == '12345A'
+    assert len(registry.devices) == 1

tests/helpers/test_area_registry.py

@@ -0,0 +1,127 @@
+"""Tests for the Area Registry."""
+import pytest
+
+from homeassistant.helpers import area_registry
+from tests.common import mock_area_registry, flush_store
+
+
+@pytest.fixture
+def registry(hass):
+    """Return an empty, loaded, registry."""
+    return mock_area_registry(hass)
+
+
+async def test_list_areas(registry):
+    """Make sure that we can read areas."""
+    registry.async_create('mock')
+
+    areas = registry.async_list_areas()
+
+    assert len(areas) == len(registry.areas)
+
+
+async def test_create_area(registry):
+    """Make sure that we can create an area."""
+    area = registry.async_create('mock')
+
+    assert area.name == 'mock'
+    assert len(registry.areas) == 1
+
+
+async def test_create_area_with_name_already_in_use(registry):
+    """Make sure that we can't create an area with a name already in use."""
+    area1 = registry.async_create('mock')
+
+    with pytest.raises(ValueError) as e_info:
+        area2 = registry.async_create('mock')
+        assert area1 != area2
+        assert e_info == "Name is already in use"
+
+    assert len(registry.areas) == 1
+
+
+async def test_delete_area(registry):
+    """Make sure that we can delete an area."""
+    area = registry.async_create('mock')
+
+    await registry.async_delete(area.id)
+
+    assert not registry.areas
+
+
+async def test_delete_non_existing_area(registry):
+    """Make sure that we can't delete an area that doesn't exist."""
+    registry.async_create('mock')
+
+    with pytest.raises(KeyError):
+        await registry.async_delete('')
+
+    assert len(registry.areas) == 1
+
+
+async def test_update_area(registry):
+    """Make sure that we can read areas."""
+    area = registry.async_create('mock')
+
+    updated_area = registry.async_update(area.id, name='mock1')
+
+    assert updated_area != area
+    assert updated_area.name == 'mock1'
+    assert len(registry.areas) == 1
+
+
+async def test_update_area_with_same_name(registry):
+    """Make sure that we can reapply the same name to the area."""
+    area = registry.async_create('mock')
+
+    updated_area = registry.async_update(area.id, name='mock')
+
+    assert updated_area == area
+    assert len(registry.areas) == 1
+
+
+async def test_update_area_with_name_already_in_use(registry):
+    """Make sure that we can't update an area with a name already in use."""
+    area1 = registry.async_create('mock1')
+    area2 = registry.async_create('mock2')
+
+    with pytest.raises(ValueError) as e_info:
+        registry.async_update(area1.id, name='mock2')
+        assert e_info == "Name is already in use"
+
+    assert area1.name == 'mock1'
+    assert area2.name == 'mock2'
+    assert len(registry.areas) == 2
+
+
+async def test_load_area(hass, registry):
+    """Make sure that we can load/save data correctly."""
+    registry.async_create('mock1')
+    registry.async_create('mock2')
+
+    assert len(registry.areas) == 2
+
+    registry2 = area_registry.AreaRegistry(hass)
+    await flush_store(registry._store)
+    await registry2.async_load()
+
+    assert list(registry.areas) == list(registry2.areas)
+
+
+async def test_loading_area_from_storage(hass, hass_storage):
+    """Test loading stored areas on start."""
+    hass_storage[area_registry.STORAGE_KEY] = {
+        'version': area_registry.STORAGE_VERSION,
+        'data': {
+            'areas': [
+                {
+                    'id': '12345A',
+                    'name': 'mock'
+                }
+            ]
+        }
+    }
+
+    registry = await area_registry.async_get_registry(hass)
+
+    assert len(registry.areas) == 1

tests/helpers/test_device_registry.py

@@ -133,6 +133,7 @@ async def test_loading_from_storage(hass, hass_storage):
                     'model': 'model',
                     'name': 'name',
                     'sw_version': 'version',
+                    'area_id': '12345A'
                 }
             ]
         }
@@ -146,6 +147,7 @@ async def test_loading_from_storage(hass, hass_storage):
         identifiers={('serial', '12:34:56:AB:CD:EF')},
         manufacturer='manufacturer', model='model')
     assert entry.id == 'abcdefghijklm'
+    assert entry.area_id == '12345A'
     assert isinstance(entry.config_entries, set)
 
 
@@ -186,6 +188,25 @@ async def test_removing_config_entries(registry):
     assert entry3.config_entries == set()
 
 
+async def test_removing_area_id(registry):
+    """Make sure we can clear area id."""
+    entry = registry.async_get_or_create(
+        config_entry_id='123',
+        connections={
+            (device_registry.CONNECTION_NETWORK_MAC, '12:34:56:AB:CD:EF')
+        },
+        identifiers={('bridgeid', '0123')},
+        manufacturer='manufacturer', model='model')
+
+    entry_w_area = registry.async_update_device(entry.id, area_id='12345A')
+
+    registry.async_clear_area_id('12345A')
+    entry_wo_area = registry.async_get_device({('bridgeid', '0123')}, set())
+
+    assert not entry_wo_area.area_id
+    assert entry_w_area != entry_wo_area
+
+
 async def test_specifying_hub_device_create(registry):
     """Test specifying a hub and updating."""
     hub = registry.async_get_or_create(
@@ -328,3 +349,19 @@ async def test_format_mac(registry):
             },
         )
         assert list(invalid_mac_entry.connections)[0][1] == invalid
+
+
+async def test_update(registry):
+    """Verify that we can update area_id of a device."""
+    entry = registry.async_get_or_create(
+        config_entry_id='1234',
+        connections={
+            (device_registry.CONNECTION_NETWORK_MAC, '12:34:56:AB:CD:EF')
+        })
+
+    assert not entry.area_id
+
+    updated_entry = registry.async_update_device(entry.id, area_id='12345A')
+
+    assert updated_entry != entry
+    assert updated_entry.area_id == '12345A'