Add support for entity categories to MQTT entities (#57656)

* Add support for entity categories to MQTT entities

* Improve test

* Apply suggestions from code review

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

* Update homeassistant/components/mqtt/mixins.py

Co-authored-by: Paul Monigatti <paulmonigatti@users.noreply.github.com>

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: Paul Monigatti <paulmonigatti@users.noreply.github.com>

homeassistant/components/mqtt/mixins.py

@@ -8,14 +8,20 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.const import CONF_DEVICE, CONF_ICON, CONF_NAME, CONF_UNIQUE_ID
+from homeassistant.const import (
+    CONF_DEVICE,
+    CONF_ENTITY_CATEGORY,
+    CONF_ICON,
+    CONF_NAME,
+    CONF_UNIQUE_ID,
+)
 from homeassistant.core import callback
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
 )
-from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.entity import ENTITY_CATEGORIES_SCHEMA, Entity
 from homeassistant.helpers.typing import ConfigType
 
 from . import DATA_MQTT, debug_info, publish, subscription
@@ -76,6 +82,7 @@ MQTT_ATTRIBUTES_BLOCKED = {
     "context_recent_time",
     "device_class",
     "device_info",
+    "entity_category",
     "entity_picture",
     "entity_registry_enabled_default",
     "extra_state_attributes",
@@ -165,6 +172,7 @@ MQTT_ENTITY_COMMON_SCHEMA = MQTT_AVAILABILITY_SCHEMA.extend(
     {
         vol.Optional(CONF_DEVICE): MQTT_ENTITY_DEVICE_INFO_SCHEMA,
         vol.Optional(CONF_ENABLED_BY_DEFAULT, default=True): cv.boolean,
+        vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA,
         vol.Optional(CONF_ICON): cv.icon,
         vol.Optional(CONF_JSON_ATTRS_TOPIC): valid_subscribe_topic,
         vol.Optional(CONF_JSON_ATTRS_TEMPLATE): cv.template,
@@ -630,6 +638,11 @@ class MqttEntity(
         """Return if the entity should be enabled when first added to the entity registry."""
         return self._config[CONF_ENABLED_BY_DEFAULT]
 
+    @property
+    def entity_category(self) -> str | None:
+        """Return the entity category if any."""
+        return self._config.get(CONF_ENTITY_CATEGORY)
+
     @property
     def icon(self):
         """Return icon of the entity if any."""

homeassistant/const.py

@@ -102,6 +102,7 @@ CONF_EFFECT: Final = "effect"
 CONF_ELEVATION: Final = "elevation"
 CONF_EMAIL: Final = "email"
 CONF_ENTITIES: Final = "entities"
+CONF_ENTITY_CATEGORY: Final = "entity_category"
 CONF_ENTITY_ID: Final = "entity_id"
 CONF_ENTITY_NAMESPACE: Final = "entity_namespace"
 CONF_ENTITY_PICTURE_TEMPLATE: Final = "entity_picture_template"

homeassistant/helpers/entity.py

@@ -11,7 +11,9 @@ import logging
 import math
 import sys
 from timeit import default_timer as timer
-from typing import Any, Literal, TypedDict, final
+from typing import Any, Final, Literal, TypedDict, final
+
+import voluptuous as vol
 
 from homeassistant.config import DATA_CUSTOMIZE
 from homeassistant.const import (
@@ -24,6 +26,8 @@ from homeassistant.const import (
     ATTR_SUPPORTED_FEATURES,
     ATTR_UNIT_OF_MEASUREMENT,
     DEVICE_DEFAULT_NAME,
+    ENTITY_CATEGORY_CONFIG,
+    ENTITY_CATEGORY_DIAGNOSTIC,
     STATE_OFF,
     STATE_ON,
     STATE_UNAVAILABLE,
@@ -52,6 +56,14 @@ SOURCE_PLATFORM_CONFIG = "platform_config"
 FLOAT_PRECISION = abs(int(math.floor(math.log10(abs(sys.float_info.epsilon))))) - 1
 
 
+ENTITY_CATEGORIES: Final[list[str]] = [
+    ENTITY_CATEGORY_CONFIG,
+    ENTITY_CATEGORY_DIAGNOSTIC,
+]
+
+ENTITY_CATEGORIES_SCHEMA: Final = vol.In(ENTITY_CATEGORIES)
+
+
 @callback
 @bind_hass
 def entity_sources(hass: HomeAssistant) -> dict[str, dict[str, str]]:

tests/components/mqtt/test_common.py

@@ -1217,3 +1217,44 @@ async def help_test_entity_disabled_by_default(hass, mqtt_mock, domain, config):
     assert not ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, "veryunique1")
     assert not ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, "veryunique2")
     assert not dev_registry.async_get_device({("mqtt", "helloworld")})
+
+
+async def help_test_entity_category(hass, mqtt_mock, domain, config):
+    """Test device registry remove."""
+    # Add device settings to config
+    config = copy.deepcopy(config[domain])
+    config["device"] = copy.deepcopy(DEFAULT_CONFIG_DEVICE_INFO_ID)
+
+    ent_registry = er.async_get(hass)
+
+    # Discover an entity without entity category
+    unique_id = "veryunique1"
+    config["unique_id"] = unique_id
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, f"homeassistant/{domain}/{unique_id}/config", data)
+    await hass.async_block_till_done()
+    entity_id = ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, unique_id)
+    assert hass.states.get(entity_id)
+    entry = ent_registry.async_get(entity_id)
+    assert entry.entity_category is None
+
+    # Discover an entity with entity category set to "config"
+    unique_id = "veryunique2"
+    config["entity_category"] = "config"
+    config["unique_id"] = unique_id
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, f"homeassistant/{domain}/{unique_id}/config", data)
+    await hass.async_block_till_done()
+    entity_id = ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, unique_id)
+    assert hass.states.get(entity_id)
+    entry = ent_registry.async_get(entity_id)
+    assert entry.entity_category == "config"
+
+    # Discover an entity with entity category set to "no_such_category"
+    unique_id = "veryunique3"
+    config["entity_category"] = "no_such_category"
+    config["unique_id"] = unique_id
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, f"homeassistant/{domain}/{unique_id}/config", data)
+    await hass.async_block_till_done()
+    assert not ent_registry.async_get_entity_id(domain, mqtt.DOMAIN, unique_id)

tests/components/mqtt/test_sensor.py

@@ -29,6 +29,7 @@ from .test_common import (
     help_test_discovery_update_attr,
     help_test_discovery_update_availability,
     help_test_discovery_update_unchanged,
+    help_test_entity_category,
     help_test_entity_debug_info,
     help_test_entity_debug_info_max_messages,
     help_test_entity_debug_info_message,
@@ -838,6 +839,12 @@ async def test_entity_disabled_by_default(hass, mqtt_mock):
     )
 
 
+@pytest.mark.no_fail_on_log_exception
+async def test_entity_category(hass, mqtt_mock):
+    """Test entity category."""
+    await help_test_entity_category(hass, mqtt_mock, sensor.DOMAIN, DEFAULT_CONFIG)
+
+
 async def test_value_template_with_entity_id(hass, mqtt_mock):
     """Test the access to attributes in value_template via the entity_id."""
     assert await async_setup_component(