Add reorder floors and areas websocket command (home-assistant#156802)

Co-authored-by: Copilot <[email protected]>

Author: piitaya

homeassistant/components/config/area_registry.py

@@ -18,6 +18,7 @@ def async_setup(hass: HomeAssistant) -> bool:
     websocket_api.async_register_command(hass, websocket_create_area)
     websocket_api.async_register_command(hass, websocket_delete_area)
     websocket_api.async_register_command(hass, websocket_update_area)
+    websocket_api.async_register_command(hass, websocket_reorder_areas)
     return True
 
 
@@ -145,3 +146,27 @@ def websocket_update_area(
         connection.send_error(msg["id"], "invalid_info", str(err))
     else:
         connection.send_result(msg["id"], entry.json_fragment)
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "config/area_registry/reorder",
+        vol.Required("area_ids"): [str],
+    }
+)
+@websocket_api.require_admin
+@callback
+def websocket_reorder_areas(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
+    """Handle reorder areas websocket command."""
+    registry = ar.async_get(hass)
+
+    try:
+        registry.async_reorder(msg["area_ids"])
+    except ValueError as err:
+        connection.send_error(msg["id"], websocket_api.ERR_INVALID_FORMAT, str(err))
+    else:
+        connection.send_result(msg["id"])

homeassistant/components/config/floor_registry.py

@@ -18,6 +18,7 @@ def async_setup(hass: HomeAssistant) -> bool:
     websocket_api.async_register_command(hass, websocket_create_floor)
     websocket_api.async_register_command(hass, websocket_delete_floor)
     websocket_api.async_register_command(hass, websocket_update_floor)
+    websocket_api.async_register_command(hass, websocket_reorder_floors)
     return True
 
 
@@ -127,6 +128,28 @@ def websocket_update_floor(
         connection.send_result(msg["id"], _entry_dict(entry))
 
 
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "config/floor_registry/reorder",
+        vol.Required("floor_ids"): [str],
+    }
+)
+@websocket_api.require_admin
+@callback
+def websocket_reorder_floors(
+    hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Handle reorder floors websocket command."""
+    registry = fr.async_get(hass)
+
+    try:
+        registry.async_reorder(msg["floor_ids"])
+    except ValueError as err:
+        connection.send_error(msg["id"], websocket_api.ERR_INVALID_FORMAT, str(err))
+    else:
+        connection.send_result(msg["id"])
+
+
 @callback
 def _entry_dict(entry: FloorEntry) -> dict[str, Any]:
     """Convert entry to API format."""

homeassistant/helpers/area_registry.py

@@ -68,8 +68,8 @@ class AreasRegistryStoreData(TypedDict):
 class EventAreaRegistryUpdatedData(TypedDict):
     """EventAreaRegistryUpdated data."""
 
-    action: Literal["create", "remove", "update"]
-    area_id: str
+    action: Literal["create", "remove", "update", "reorder"]
+    area_id: str | None
 
 
 @dataclass(frozen=True, kw_only=True, slots=True)
@@ -420,6 +420,26 @@ def _async_update(
         self.async_schedule_save()
         return new
 
+    @callback
+    def async_reorder(self, area_ids: list[str]) -> None:
+        """Reorder areas."""
+        self.hass.verify_event_loop_thread("area_registry.async_reorder")
+
+        if set(area_ids) != set(self.areas.data.keys()):
+            raise ValueError(
+                "The area_ids list must contain all existing area IDs exactly once"
+            )
+
+        reordered_data = {area_id: self.areas.data[area_id] for area_id in area_ids}
+        self.areas.data.clear()
+        self.areas.data.update(reordered_data)
+
+        self.async_schedule_save()
+        self.hass.bus.async_fire_internal(
+            EVENT_AREA_REGISTRY_UPDATED,
+            EventAreaRegistryUpdatedData(action="reorder", area_id=None),
+        )
+
     async def async_load(self) -> None:
         """Load the area registry."""
         self._async_setup_cleanup()
@@ -489,7 +509,9 @@ def _removed_from_registry_filter(
         @callback
         def _handle_floor_registry_update(event: fr.EventFloorRegistryUpdated) -> None:
             """Update areas that are associated with a floor that has been removed."""
-            floor_id = event.data["floor_id"]
+            floor_id = event.data.get("floor_id")
+            if floor_id is None:
+                return
             for area in self.areas.get_areas_for_floor(floor_id):
                 self.async_update(area.id, floor_id=None)
 

homeassistant/helpers/floor_registry.py

@@ -54,8 +54,8 @@ class FloorRegistryStoreData(TypedDict):
 class EventFloorRegistryUpdatedData(TypedDict):
     """Event data for when the floor registry is updated."""
 
-    action: Literal["create", "remove", "update"]
-    floor_id: str
+    action: Literal["create", "remove", "update", "reorder"]
+    floor_id: str | None
 
 
 type EventFloorRegistryUpdated = Event[EventFloorRegistryUpdatedData]
@@ -261,6 +261,28 @@ def async_update(
 
         return new
 
+    @callback
+    def async_reorder(self, floor_ids: list[str]) -> None:
+        """Reorder floors."""
+        self.hass.verify_event_loop_thread("floor_registry.async_reorder")
+
+        if set(floor_ids) != set(self.floors.data.keys()):
+            raise ValueError(
+                "The floor_ids list must contain all existing floor IDs exactly once"
+            )
+
+        reordered_data = {
+            floor_id: self.floors.data[floor_id] for floor_id in floor_ids
+        }
+        self.floors.data.clear()
+        self.floors.data.update(reordered_data)
+
+        self.async_schedule_save()
+        self.hass.bus.async_fire_internal(
+            EVENT_FLOOR_REGISTRY_UPDATED,
+            EventFloorRegistryUpdatedData(action="reorder", floor_id=None),
+        )
+
     async def async_load(self) -> None:
         """Load the floor registry."""
         data = await self._store.async_load()

tests/components/config/test_area_registry.py

@@ -1,6 +1,7 @@
 """Test area_registry API."""
 
 from datetime import datetime
+from typing import Any
 
 from freezegun.api import FrozenDateTimeFactory
 import pytest
@@ -346,3 +347,92 @@ async def test_update_area_with_name_already_in_use(
     assert msg["error"]["code"] == "invalid_info"
     assert msg["error"]["message"] == "The name mock 2 (mock2) is already in use"
     assert len(area_registry.areas) == 2
+
+
+async def test_reorder_areas(
+    client: MockHAClientWebSocket, area_registry: ar.AreaRegistry
+) -> None:
+    """Test reorder areas."""
+    area1 = area_registry.async_create("mock 1")
+    area2 = area_registry.async_create("mock 2")
+    area3 = area_registry.async_create("mock 3")
+
+    await client.send_json_auto_id({"type": "config/area_registry/list"})
+    msg = await client.receive_json()
+    assert [area["area_id"] for area in msg["result"]] == [area1.id, area2.id, area3.id]
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/area_registry/reorder",
+            "area_ids": [area3.id, area1.id, area2.id],
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+
+    await client.send_json_auto_id({"type": "config/area_registry/list"})
+    msg = await client.receive_json()
+    assert [area["area_id"] for area in msg["result"]] == [area3.id, area1.id, area2.id]
+
+
+async def test_reorder_areas_invalid_area_ids(
+    client: MockHAClientWebSocket, area_registry: ar.AreaRegistry
+) -> None:
+    """Test reorder with invalid area IDs."""
+    area1 = area_registry.async_create("mock 1")
+    area_registry.async_create("mock 2")
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/area_registry/reorder",
+            "area_ids": [area1.id],
+        }
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == "invalid_format"
+    assert "must contain all existing area IDs" in msg["error"]["message"]
+
+
+async def test_reorder_areas_with_nonexistent_id(
+    client: MockHAClientWebSocket, area_registry: ar.AreaRegistry
+) -> None:
+    """Test reorder with nonexistent area ID."""
+    area1 = area_registry.async_create("mock 1")
+    area2 = area_registry.async_create("mock 2")
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/area_registry/reorder",
+            "area_ids": [area1.id, area2.id, "nonexistent"],
+        }
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == "invalid_format"
+
+
+async def test_reorder_areas_persistence(
+    hass: HomeAssistant,
+    client: MockHAClientWebSocket,
+    area_registry: ar.AreaRegistry,
+    hass_storage: dict[str, Any],
+) -> None:
+    """Test that area reordering is persisted."""
+    area1 = area_registry.async_create("mock 1")
+    area2 = area_registry.async_create("mock 2")
+    area3 = area_registry.async_create("mock 3")
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/area_registry/reorder",
+            "area_ids": [area2.id, area3.id, area1.id],
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+
+    await hass.async_block_till_done()
+
+    area_ids = [area.id for area in area_registry.async_list_areas()]
+    assert area_ids == [area2.id, area3.id, area1.id]

tests/components/config/test_floor_registry.py

@@ -1,6 +1,7 @@
 """Test floor registry API."""
 
 from datetime import datetime
+from typing import Any
 
 from freezegun.api import FrozenDateTimeFactory
 import pytest
@@ -275,3 +276,100 @@ async def test_update_with_name_already_in_use(
         == "The name Second floor (secondfloor) is already in use"
     )
     assert len(floor_registry.floors) == 2
+
+
+async def test_reorder_floors(
+    client: MockHAClientWebSocket, floor_registry: fr.FloorRegistry
+) -> None:
+    """Test reorder floors."""
+    floor1 = floor_registry.async_create("First floor")
+    floor2 = floor_registry.async_create("Second floor")
+    floor3 = floor_registry.async_create("Third floor")
+
+    await client.send_json_auto_id({"type": "config/floor_registry/list"})
+    msg = await client.receive_json()
+    assert [floor["floor_id"] for floor in msg["result"]] == [
+        floor1.floor_id,
+        floor2.floor_id,
+        floor3.floor_id,
+    ]
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/floor_registry/reorder",
+            "floor_ids": [floor3.floor_id, floor1.floor_id, floor2.floor_id],
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+
+    await client.send_json_auto_id({"type": "config/floor_registry/list"})
+    msg = await client.receive_json()
+    assert [floor["floor_id"] for floor in msg["result"]] == [
+        floor3.floor_id,
+        floor1.floor_id,
+        floor2.floor_id,
+    ]
+
+
+async def test_reorder_floors_invalid_floor_ids(
+    client: MockHAClientWebSocket, floor_registry: fr.FloorRegistry
+) -> None:
+    """Test reorder with invalid floor IDs."""
+    floor1 = floor_registry.async_create("First floor")
+    floor_registry.async_create("Second floor")
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/floor_registry/reorder",
+            "floor_ids": [floor1.floor_id],
+        }
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == "invalid_format"
+    assert "must contain all existing floor IDs" in msg["error"]["message"]
+
+
+async def test_reorder_floors_with_nonexistent_id(
+    client: MockHAClientWebSocket, floor_registry: fr.FloorRegistry
+) -> None:
+    """Test reorder with nonexistent floor ID."""
+    floor1 = floor_registry.async_create("First floor")
+    floor2 = floor_registry.async_create("Second floor")
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/floor_registry/reorder",
+            "floor_ids": [floor1.floor_id, floor2.floor_id, "nonexistent"],
+        }
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == "invalid_format"
+
+
+async def test_reorder_floors_persistence(
+    hass: HomeAssistant,
+    client: MockHAClientWebSocket,
+    floor_registry: fr.FloorRegistry,
+    hass_storage: dict[str, Any],
+) -> None:
+    """Test that floor reordering is persisted."""
+    floor1 = floor_registry.async_create("First floor")
+    floor2 = floor_registry.async_create("Second floor")
+    floor3 = floor_registry.async_create("Third floor")
+
+    await client.send_json_auto_id(
+        {
+            "type": "config/floor_registry/reorder",
+            "floor_ids": [floor2.floor_id, floor3.floor_id, floor1.floor_id],
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+
+    await hass.async_block_till_done()
+
+    floor_ids = [floor.floor_id for floor in floor_registry.async_list_floors()]
+    assert floor_ids == [floor2.floor_id, floor3.floor_id, floor1.floor_id]