From d5443b8dee81c51840a36aaee43775221f84628c Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 8 Feb 2022 23:00:53 +0100 Subject: [PATCH] Fix ENTITY_CATEGORIES_SCHEMA (#66108) Co-authored-by: Paulus Schoutsen --- homeassistant/components/knx/schema.py | 26 +++++++++---------- .../components/mobile_app/webhook.py | 4 +-- homeassistant/components/mqtt/mixins.py | 4 +-- homeassistant/helpers/entity.py | 10 +++++-- tests/helpers/test_entity.py | 25 ++++++++++++++++++ 5 files changed, 50 insertions(+), 19 deletions(-) diff --git a/homeassistant/components/knx/schema.py b/homeassistant/components/knx/schema.py index af475e9c380..e946266f8f4 100644 --- a/homeassistant/components/knx/schema.py +++ b/homeassistant/components/knx/schema.py @@ -35,7 +35,7 @@ from homeassistant.const import ( Platform, ) import homeassistant.helpers.config_validation as cv -from homeassistant.helpers.entity import ENTITY_CATEGORIES_SCHEMA +from homeassistant.helpers.entity import validate_entity_category from .const import ( CONF_INVERT, @@ -320,7 +320,7 @@ class BinarySensorSchema(KNXPlatformSchema): ), vol.Optional(CONF_DEVICE_CLASS): BINARY_SENSOR_DEVICE_CLASSES_SCHEMA, vol.Optional(CONF_RESET_AFTER): cv.positive_float, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), ) @@ -356,7 +356,7 @@ class ButtonSchema(KNXPlatformSchema): vol.Exclusive( CONF_TYPE, "length_or_type", msg=length_or_type_msg ): object, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), vol.Any( @@ -500,7 +500,7 @@ class ClimateSchema(KNXPlatformSchema): ): vol.In(HVAC_MODES), vol.Optional(CONF_MIN_TEMP): vol.Coerce(float), vol.Optional(CONF_MAX_TEMP): vol.Coerce(float), - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), ) @@ -555,7 +555,7 @@ class CoverSchema(KNXPlatformSchema): vol.Optional(CONF_INVERT_POSITION, default=False): cv.boolean, vol.Optional(CONF_INVERT_ANGLE, default=False): cv.boolean, vol.Optional(CONF_DEVICE_CLASS): COVER_DEVICE_CLASSES_SCHEMA, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), ) @@ -618,7 +618,7 @@ class FanSchema(KNXPlatformSchema): vol.Optional(CONF_OSCILLATION_ADDRESS): ga_list_validator, vol.Optional(CONF_OSCILLATION_STATE_ADDRESS): ga_list_validator, vol.Optional(CONF_MAX_STEP): cv.byte, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ) @@ -722,7 +722,7 @@ class LightSchema(KNXPlatformSchema): vol.Optional(CONF_MAX_KELVIN, default=DEFAULT_MAX_KELVIN): vol.All( vol.Coerce(int), vol.Range(min=1) ), - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), vol.Any( @@ -802,7 +802,7 @@ class NumberSchema(KNXPlatformSchema): vol.Optional(CONF_MAX): vol.Coerce(float), vol.Optional(CONF_MIN): vol.Coerce(float), vol.Optional(CONF_STEP): cv.positive_float, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), number_limit_sub_validator, @@ -824,7 +824,7 @@ class SceneSchema(KNXPlatformSchema): vol.Required(CONF_SCENE_NUMBER): vol.All( vol.Coerce(int), vol.Range(min=1, max=64) ), - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ) @@ -855,7 +855,7 @@ class SelectSchema(KNXPlatformSchema): ], vol.Required(KNX_ADDRESS): ga_list_validator, vol.Optional(CONF_STATE_ADDRESS): ga_list_validator, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), select_options_sub_validator, @@ -880,7 +880,7 @@ class SensorSchema(KNXPlatformSchema): vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA, vol.Required(CONF_TYPE): sensor_type_validator, vol.Required(CONF_STATE_ADDRESS): ga_list_validator, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ) @@ -901,7 +901,7 @@ class SwitchSchema(KNXPlatformSchema): vol.Optional(CONF_RESPOND_TO_READ, default=False): cv.boolean, vol.Required(KNX_ADDRESS): ga_list_validator, vol.Optional(CONF_STATE_ADDRESS): ga_list_validator, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ) @@ -948,7 +948,7 @@ class WeatherSchema(KNXPlatformSchema): vol.Optional(CONF_KNX_DAY_NIGHT_ADDRESS): ga_list_validator, vol.Optional(CONF_KNX_AIR_PRESSURE_ADDRESS): ga_list_validator, vol.Optional(CONF_KNX_HUMIDITY_ADDRESS): ga_list_validator, - vol.Optional(CONF_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(CONF_ENTITY_CATEGORY): validate_entity_category, } ), ) diff --git a/homeassistant/components/mobile_app/webhook.py b/homeassistant/components/mobile_app/webhook.py index d093fd23406..d659d7625c1 100644 --- a/homeassistant/components/mobile_app/webhook.py +++ b/homeassistant/components/mobile_app/webhook.py @@ -44,7 +44,7 @@ from homeassistant.helpers import ( template, ) from homeassistant.helpers.dispatcher import async_dispatcher_send -from homeassistant.helpers.entity import ENTITY_CATEGORIES_SCHEMA +from homeassistant.helpers.entity import validate_entity_category from homeassistant.util.decorator import Registry from .const import ( @@ -423,7 +423,7 @@ def _validate_state_class_sensor(value: dict): vol.Optional(ATTR_SENSOR_STATE, default=None): vol.Any( None, bool, str, int, float ), - vol.Optional(ATTR_SENSOR_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(ATTR_SENSOR_ENTITY_CATEGORY): validate_entity_category, vol.Optional(ATTR_SENSOR_ICON, default="mdi:cellphone"): cv.icon, vol.Optional(ATTR_SENSOR_STATE_CLASS): vol.In(SENSOSR_STATE_CLASSES), }, diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py index 47fc358dfdc..6b92ab91e31 100644 --- a/homeassistant/components/mqtt/mixins.py +++ b/homeassistant/components/mqtt/mixins.py @@ -30,11 +30,11 @@ from homeassistant.helpers.dispatcher import ( async_dispatcher_send, ) from homeassistant.helpers.entity import ( - ENTITY_CATEGORIES_SCHEMA, DeviceInfo, Entity, EntityCategory, async_generate_entity_id, + validate_entity_category, ) from homeassistant.helpers.typing import ConfigType @@ -191,7 +191,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_ENTITY_CATEGORY): validate_entity_category, vol.Optional(CONF_ICON): cv.icon, vol.Optional(CONF_JSON_ATTRS_TOPIC): valid_subscribe_topic, vol.Optional(CONF_JSON_ATTRS_TEMPLATE): cv.template, diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index a716e465450..b670c734b47 100644 --- a/homeassistant/helpers/entity.py +++ b/homeassistant/helpers/entity.py @@ -11,7 +11,7 @@ import logging import math import sys from timeit import default_timer as timer -from typing import Any, Final, Literal, TypedDict, final +from typing import Any, Literal, TypedDict, final import voluptuous as vol @@ -58,7 +58,13 @@ SOURCE_PLATFORM_CONFIG = "platform_config" FLOAT_PRECISION = abs(int(math.floor(math.log10(abs(sys.float_info.epsilon))))) - 1 -ENTITY_CATEGORIES_SCHEMA: Final = vol.In(ENTITY_CATEGORIES) +def validate_entity_category(value: Any | None) -> EntityCategory: + """Validate entity category configuration.""" + value = vol.In(ENTITY_CATEGORIES)(value) + return EntityCategory(value) + + +ENTITY_CATEGORIES_SCHEMA = validate_entity_category @callback diff --git a/tests/helpers/test_entity.py b/tests/helpers/test_entity.py index 1d28c50b9a1..6b7de074a24 100644 --- a/tests/helpers/test_entity.py +++ b/tests/helpers/test_entity.py @@ -6,6 +6,7 @@ import threading from unittest.mock import MagicMock, PropertyMock, patch import pytest +import voluptuous as vol from homeassistant.const import ( ATTR_ATTRIBUTION, @@ -829,3 +830,27 @@ async def test_entity_category_property(hass): ) mock_entity2.entity_id = "hello.world" assert mock_entity2.entity_category == "config" + + +@pytest.mark.parametrize( + "value,expected", + ( + ("config", entity.EntityCategory.CONFIG), + ("diagnostic", entity.EntityCategory.DIAGNOSTIC), + ("system", entity.EntityCategory.SYSTEM), + ), +) +def test_entity_category_schema(value, expected): + """Test entity category schema.""" + schema = vol.Schema(entity.ENTITY_CATEGORIES_SCHEMA) + result = schema(value) + assert result == expected + assert isinstance(result, entity.EntityCategory) + + +@pytest.mark.parametrize("value", (None, "non_existing")) +def test_entity_category_schema_error(value): + """Test entity category schema.""" + schema = vol.Schema(entity.ENTITY_CATEGORIES_SCHEMA) + with pytest.raises(vol.Invalid): + schema(value)