Add floor registry (#110741)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

tests/helpers/test_area_registry.py

@@ -4,7 +4,7 @@ from typing import Any
 import pytest
 
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import area_registry as ar
+from homeassistant.helpers import area_registry as ar, floor_registry as fr
 
 from tests.common import ANY, async_capture_events, flush_store
 
@@ -27,6 +27,7 @@ async def test_create_area(hass: HomeAssistant, area_registry: ar.AreaRegistry)
 
     assert area == ar.AreaEntry(
         aliases=set(),
+        floor_id=None,
         icon=None,
         id=ANY,
         name="mock",
@@ -50,6 +51,7 @@ async def test_create_area(hass: HomeAssistant, area_registry: ar.AreaRegistry)
 
     assert area == ar.AreaEntry(
         aliases={"alias_1", "alias_2"},
+        floor_id=None,
         icon=None,
         id=ANY,
         name="mock 2",
@@ -133,14 +135,20 @@ async def test_delete_non_existing_area(area_registry: ar.AreaRegistry) -> None:
     assert len(area_registry.areas) == 1
 
 
-async def test_update_area(hass: HomeAssistant, area_registry: ar.AreaRegistry) -> None:
+async def test_update_area(
+    hass: HomeAssistant,
+    area_registry: ar.AreaRegistry,
+    floor_registry: fr.FloorRegistry,
+) -> None:
     """Make sure that we can read areas."""
     update_events = async_capture_events(hass, ar.EVENT_AREA_REGISTRY_UPDATED)
+    floor_registry.async_create("first")
     area = area_registry.async_create("mock")
 
     updated_area = area_registry.async_update(
         area.id,
         aliases={"alias_1", "alias_2"},
+        floor_id="first",
         icon="mdi:garage",
         name="mock1",
         picture="/image/example.png",
@@ -149,6 +157,7 @@ async def test_update_area(hass: HomeAssistant, area_registry: ar.AreaRegistry)
     assert updated_area != area
     assert updated_area == ar.AreaEntry(
         aliases={"alias_1", "alias_2"},
+        floor_id="first",
         icon="mdi:garage",
         id=ANY,
         name="mock1",
@@ -257,6 +266,7 @@ async def test_loading_area_from_storage(
             "areas": [
                 {
                     "aliases": ["alias_1", "alias_2"],
+                    "floor_id": "first_floor",
                     "id": "12345A",
                     "icon": "mdi:garage",
                     "name": "mock",
@@ -299,6 +309,7 @@ async def test_migration_from_1_1(
             "areas": [
                 {
                     "aliases": [],
+                    "floor_id": None,
                     "icon": None,
                     "id": "12345A",
                     "name": "mock",
@@ -345,3 +356,58 @@ async def test_async_get_area(area_registry: ar.AreaRegistry) -> None:
     assert len(area_registry.areas) == 1
 
     assert area_registry.async_get_area(area.id).normalized_name == "mock1"
+
+
+async def test_removing_floors(
+    hass: HomeAssistant,
+    area_registry: ar.AreaRegistry,
+    floor_registry: fr.FloorRegistry,
+) -> None:
+    """Make sure we can clear floors."""
+    first_floor = floor_registry.async_create("First floor")
+    second_floor = floor_registry.async_create("Second floor")
+
+    kitchen = area_registry.async_create("Kitchen")
+    kitchen = area_registry.async_update(kitchen.id, floor_id=first_floor.floor_id)
+    bedroom = area_registry.async_create("Bedroom")
+    bedroom = area_registry.async_update(bedroom.id, floor_id=second_floor.floor_id)
+
+    floor_registry.async_delete(first_floor.floor_id)
+    await hass.async_block_till_done()
+    assert area_registry.async_get_area(kitchen.id).floor_id is None
+    assert area_registry.async_get_area(bedroom.id).floor_id == second_floor.floor_id
+
+    floor_registry.async_delete(second_floor.floor_id)
+    await hass.async_block_till_done()
+    assert area_registry.async_get_area(kitchen.id).floor_id is None
+    assert area_registry.async_get_area(bedroom.id).floor_id is None
+
+
+@pytest.mark.usefixtures("hass")
+async def test_entries_for_floor(
+    area_registry: ar.AreaRegistry,
+    floor_registry: fr.FloorRegistry,
+) -> None:
+    """Test getting area entries by floor."""
+    first_floor = floor_registry.async_create("First floor")
+    second_floor = floor_registry.async_create("Second floor")
+
+    kitchen = area_registry.async_create("Kitchen")
+    kitchen = area_registry.async_update(kitchen.id, floor_id=first_floor.floor_id)
+    living_room = area_registry.async_create("Living room")
+    living_room = area_registry.async_update(
+        living_room.id, floor_id=first_floor.floor_id
+    )
+    bedroom = area_registry.async_create("Bedroom")
+    bedroom = area_registry.async_update(bedroom.id, floor_id=second_floor.floor_id)
+
+    entries = ar.async_entries_for_floor(area_registry, first_floor.floor_id)
+    assert len(entries) == 2
+    assert entries == [kitchen, living_room]
+
+    entries = ar.async_entries_for_floor(area_registry, second_floor.floor_id)
+    assert len(entries) == 1
+    assert entries == [bedroom]
+
+    assert not ar.async_entries_for_floor(area_registry, "unknown")
+    assert not ar.async_entries_for_floor(area_registry, "")