From 1eebe451544f95c2fcb395055fdc2d736617885e Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Fri, 15 Oct 2021 14:28:30 +0200 Subject: [PATCH] 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 * Update homeassistant/components/mqtt/mixins.py Co-authored-by: Paul Monigatti Co-authored-by: Paulus Schoutsen Co-authored-by: Paul Monigatti --- homeassistant/components/mqtt/mixins.py | 17 ++++++++-- homeassistant/const.py | 1 + homeassistant/helpers/entity.py | 14 ++++++++- tests/components/mqtt/test_common.py | 41 +++++++++++++++++++++++++ tests/components/mqtt/test_sensor.py | 7 +++++ 5 files changed, 77 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py index 0ba699c6229b..1965cb77e531 100644 --- a/homeassistant/components/mqtt/mixins.py +++ b/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.""" diff --git a/homeassistant/const.py b/homeassistant/const.py index f4387c9bbce7..1c3029a9d342 100644 --- a/homeassistant/const.py +++ b/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" diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index 2ad0f934973e..a05d2c7c2fa5 100644 --- a/homeassistant/helpers/entity.py +++ b/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]]: diff --git a/tests/components/mqtt/test_common.py b/tests/components/mqtt/test_common.py index 0458ad14d51a..72f5f236d280 100644 --- a/tests/components/mqtt/test_common.py +++ b/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) diff --git a/tests/components/mqtt/test_sensor.py b/tests/components/mqtt/test_sensor.py index 46c06f0d3b39..b7120b99f6ee 100644 --- a/tests/components/mqtt/test_sensor.py +++ b/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(