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

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'