From c28b45cd83d467daad0598178958fb532cb18385 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Wed, 1 Dec 2021 13:54:36 +0100
Subject: [PATCH] Migrate entity categories to StrEnum (#60720)

---
 homeassistant/const.py          |  9 ++-------
 homeassistant/helpers/entity.py | 27 ++++++++++++++++++++++++---
 tests/helpers/test_entity.py    |  4 ++--
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/homeassistant/const.py b/homeassistant/const.py
index 2ed777c12af..fb8dfc6089a 100644
--- a/homeassistant/const.py
+++ b/homeassistant/const.py
@@ -703,16 +703,11 @@ PRECISION_TENTHS: Final = 0.1
 # cloud, alexa, or google_home components
 CLOUD_NEVER_EXPOSED_ENTITIES: Final[list[str]] = ["group.all_locks"]
 
-# Config: An entity which allows changing the configuration of a device
+# ENTITY_CATEGOR* below are deprecated as of 2021.12
+# use the EntityCategory enum instead.
 ENTITY_CATEGORY_CONFIG: Final = "config"
-# Diagnostic: An entity exposing some configuration parameter or diagnostics of a device
 ENTITY_CATEGORY_DIAGNOSTIC: Final = "diagnostic"
-# System: An entity which is not useful for the user to interact with
 ENTITY_CATEGORY_SYSTEM: Final = "system"
-
-# Entity categories which will:
-# - Not be exposed to cloud, alexa, or google_home components
-# - Not be included in indirect service calls to devices or areas
 ENTITY_CATEGORIES: Final[list[str]] = [
     ENTITY_CATEGORY_CONFIG,
     ENTITY_CATEGORY_DIAGNOSTIC,
diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py
index 4d301e7d1b6..c303b97eb5f 100644
--- a/homeassistant/helpers/entity.py
+++ b/homeassistant/helpers/entity.py
@@ -43,6 +43,7 @@ from homeassistant.helpers.event import Event, async_track_entity_registry_updat
 from homeassistant.helpers.typing import StateType
 from homeassistant.loader import bind_hass
 from homeassistant.util import dt as dt_util, ensure_unique_string, slugify
+from homeassistant.util.enum import StrEnum
 
 _LOGGER = logging.getLogger(__name__)
 SLOW_UPDATE_WARNING = 10
@@ -179,6 +180,24 @@ class DeviceInfo(TypedDict, total=False):
     via_device: tuple[str, str]
 
 
+class EntityCategory(StrEnum):
+    """Category of an entity.
+
+    An entity with a category will:
+    - Not be exposed to cloud, Alexa, or Google Assistant components
+    - Not be included in indirect service calls to devices or areas
+    """
+
+    # Config: An entity which allows changing the configuration of a device
+    CONFIG = "config"
+
+    # Diagnostic: An entity exposing some configuration parameter or diagnostics of a device
+    DIAGNOSTIC = "diagnostic"
+
+    # System: An entity which is not useful for the user to interact with
+    SYSTEM = "system"
+
+
 @dataclass
 class EntityDescription:
     """A class that describes Home Assistant entities."""
@@ -187,7 +206,9 @@ class EntityDescription:
     key: str
 
     device_class: str | None = None
-    entity_category: Literal["config", "diagnostic", "system"] | None = None
+    entity_category: EntityCategory | Literal[
+        "config", "diagnostic", "system"
+    ] | None = None
     entity_registry_enabled_default: bool = True
     force_update: bool = False
     icon: str | None = None
@@ -246,7 +267,7 @@ class Entity(ABC):
     _attr_context_recent_time: timedelta = timedelta(seconds=5)
     _attr_device_class: str | None
     _attr_device_info: DeviceInfo | None = None
-    _attr_entity_category: str | None
+    _attr_entity_category: EntityCategory | str | None
     _attr_entity_picture: str | None = None
     _attr_entity_registry_enabled_default: bool
     _attr_extra_state_attributes: MutableMapping[str, Any]
@@ -414,7 +435,7 @@ class Entity(ABC):
         return self._attr_attribution
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | str | None:
         """Return the category of the entity, if any."""
         if hasattr(self, "_attr_entity_category"):
             return self._attr_entity_category
diff --git a/tests/helpers/test_entity.py b/tests/helpers/test_entity.py
index e334e6d2c56..d4051613c03 100644
--- a/tests/helpers/test_entity.py
+++ b/tests/helpers/test_entity.py
@@ -819,13 +819,13 @@ async def test_entity_category_property(hass):
         key="abc", entity_category="ignore_me"
     )
     mock_entity1.entity_id = "hello.world"
-    mock_entity1._attr_entity_category = "config"
+    mock_entity1._attr_entity_category = entity.EntityCategory.CONFIG
     assert mock_entity1.entity_category == "config"
 
     mock_entity2 = entity.Entity()
     mock_entity2.hass = hass
     mock_entity2.entity_description = entity.EntityDescription(
-        key="abc", entity_category="config"
+        key="abc", entity_category=entity.EntityCategory.CONFIG
     )
     mock_entity2.entity_id = "hello.world"
     assert mock_entity2.entity_category == "config"