From 541d4b78ac6e9009f3d4b67dec13305ed0578495 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Mon, 18 Mar 2024 22:17:13 +0100
Subject: [PATCH] Add WebSocket support for handling labels on areas registry
 (#113755)

---
 homeassistant/components/config/area_registry.py | 11 +++++++++++
 tests/components/config/test_area_registry.py    | 10 ++++++++++
 2 files changed, 21 insertions(+)

diff --git a/homeassistant/components/config/area_registry.py b/homeassistant/components/config/area_registry.py
index acddf7bb200..a499ab84784 100644
--- a/homeassistant/components/config/area_registry.py
+++ b/homeassistant/components/config/area_registry.py
@@ -42,6 +42,7 @@ def websocket_list_areas(
         vol.Optional("aliases"): list,
         vol.Optional("floor_id"): str,
         vol.Optional("icon"): str,
+        vol.Optional("labels"): [str],
         vol.Required("name"): str,
         vol.Optional("picture"): vol.Any(str, None),
     }
@@ -64,6 +65,10 @@ def websocket_create_area(
         # Convert aliases to a set
         data["aliases"] = set(data["aliases"])
 
+    if "labels" in data:
+        # Convert labels to a set
+        data["labels"] = set(data["labels"])
+
     try:
         entry = registry.async_create(**data)
     except ValueError as err:
@@ -103,6 +108,7 @@ def websocket_delete_area(
         vol.Required("area_id"): str,
         vol.Optional("floor_id"): vol.Any(str, None),
         vol.Optional("icon"): vol.Any(str, None),
+        vol.Optional("labels"): [str],
         vol.Optional("name"): str,
         vol.Optional("picture"): vol.Any(str, None),
     }
@@ -125,6 +131,10 @@ def websocket_update_area(
         # Convert aliases to a set
         data["aliases"] = set(data["aliases"])
 
+    if "labels" in data:
+        # Convert labels to a set
+        data["labels"] = set(data["labels"])
+
     try:
         entry = registry.async_update(**data)
     except ValueError as err:
@@ -141,6 +151,7 @@ def _entry_dict(entry: AreaEntry) -> dict[str, Any]:
         "area_id": entry.id,
         "floor_id": entry.floor_id,
         "icon": entry.icon,
+        "labels": list(entry.labels),
         "name": entry.name,
         "picture": entry.picture,
     }
diff --git a/tests/components/config/test_area_registry.py b/tests/components/config/test_area_registry.py
index 8bedd1a0a29..fb59725fd29 100644
--- a/tests/components/config/test_area_registry.py
+++ b/tests/components/config/test_area_registry.py
@@ -31,6 +31,7 @@ async def test_list_areas(
         icon="mdi:garage",
         picture="/image/example.png",
         floor_id="first_floor",
+        labels={"label_1", "label_2"},
     )
 
     await client.send_json_auto_id({"type": "config/area_registry/list"})
@@ -42,6 +43,7 @@ async def test_list_areas(
             "area_id": area1.id,
             "floor_id": None,
             "icon": None,
+            "labels": [],
             "name": "mock 1",
             "picture": None,
         },
@@ -50,6 +52,7 @@ async def test_list_areas(
             "area_id": area2.id,
             "floor_id": "first_floor",
             "icon": "mdi:garage",
+            "labels": unordered(["label_1", "label_2"]),
             "name": "mock 2",
             "picture": "/image/example.png",
         },
@@ -72,6 +75,7 @@ async def test_create_area(
         "area_id": ANY,
         "floor_id": None,
         "icon": None,
+        "labels": [],
         "name": "mock",
         "picture": None,
     }
@@ -83,6 +87,7 @@ async def test_create_area(
             "aliases": ["alias_1", "alias_2"],
             "floor_id": "first_floor",
             "icon": "mdi:garage",
+            "labels": ["label_1", "label_2"],
             "name": "mock 2",
             "picture": "/image/example.png",
             "type": "config/area_registry/create",
@@ -96,6 +101,7 @@ async def test_create_area(
         "area_id": ANY,
         "floor_id": "first_floor",
         "icon": "mdi:garage",
+        "labels": unordered(["label_1", "label_2"]),
         "name": "mock 2",
         "picture": "/image/example.png",
     }
@@ -166,6 +172,7 @@ async def test_update_area(
             "area_id": area.id,
             "floor_id": "first_floor",
             "icon": "mdi:garage",
+            "labels": ["label_1", "label_2"],
             "name": "mock 2",
             "picture": "/image/example.png",
             "type": "config/area_registry/update",
@@ -179,6 +186,7 @@ async def test_update_area(
         "area_id": area.id,
         "floor_id": "first_floor",
         "icon": "mdi:garage",
+        "labels": unordered(["label_1", "label_2"]),
         "name": "mock 2",
         "picture": "/image/example.png",
     }
@@ -190,6 +198,7 @@ async def test_update_area(
             "area_id": area.id,
             "floor_id": None,
             "icon": None,
+            "labels": [],
             "picture": None,
             "type": "config/area_registry/update",
         }
@@ -202,6 +211,7 @@ async def test_update_area(
         "area_id": area.id,
         "floor_id": None,
         "icon": None,
+        "labels": [],
         "name": "mock 2",
         "picture": None,
     }