Enforce EntityCategory enum (#69015)

Co-authored-by: Franck Nijhof <git@frenck.dev>
pull/69044/head
Paulus Schoutsen 2022-03-31 15:04:33 -07:00 committed by GitHub
parent 824066f519
commit 130ca2213f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 63 additions and 90 deletions

View File

@ -814,7 +814,7 @@ class MqttEntity(
return self._config[CONF_ENABLED_BY_DEFAULT]
@property
def entity_category(self) -> EntityCategory | str | None:
def entity_category(self) -> EntityCategory | None:
"""Return the entity category if any."""
return self._config.get(CONF_ENTITY_CATEGORY)

View File

@ -108,7 +108,7 @@ class NeatoConnectedSwitch(SwitchEntity):
)
@property
def entity_category(self) -> str:
def entity_category(self) -> EntityCategory:
"""Device entity category."""
return EntityCategory.CONFIG

View File

@ -183,7 +183,7 @@ class UpdateEntity(RestoreEntity):
return None
@property
def entity_category(self) -> EntityCategory | str | None:
def entity_category(self) -> EntityCategory | None:
"""Return the category of the entity, if any."""
if hasattr(self, "_attr_entity_category"):
return self._attr_entity_category

View File

@ -52,7 +52,6 @@ from . import entity_registry as er
from .device_registry import DeviceEntryType
from .entity_platform import EntityPlatform
from .event import async_track_entity_registry_updated_event
from .frame import report
from .typing import StateType
_LOGGER = logging.getLogger(__name__)
@ -228,29 +227,6 @@ class EntityPlatformState(Enum):
REMOVED = auto()
def convert_to_entity_category(
value: EntityCategory | str | None, raise_report: bool = True
) -> EntityCategory | None:
"""Force incoming entity_category to be an enum."""
if value is None:
return value
if not isinstance(value, EntityCategory):
if raise_report:
report(
"uses %s (%s) for entity category. This is deprecated and will "
"stop working in Home Assistant 2022.4, it should be updated to use "
"EntityCategory instead" % (type(value).__name__, value),
error_if_core=False,
)
try:
return EntityCategory(value)
except ValueError:
return None
return value
@dataclass
class EntityDescription:
"""A class that describes Home Assistant entities."""
@ -259,10 +235,7 @@ class EntityDescription:
key: str
device_class: str | None = None
# Type string is deprecated as of 2021.12, use EntityCategory
entity_category: EntityCategory | Literal[
"config", "diagnostic", "system"
] | None = None
entity_category: EntityCategory | None = None
entity_registry_enabled_default: bool = True
force_update: bool = False
icon: str | None = None
@ -491,9 +464,8 @@ class Entity(ABC):
"""Return the attribution."""
return self._attr_attribution
# Type str is deprecated as of 2021.12, use EntityCategory
@property
def entity_category(self) -> EntityCategory | str | None:
def entity_category(self) -> EntityCategory | None:
"""Return the category of the entity, if any."""
if hasattr(self, "_attr_entity_category"):
return self._attr_entity_category

View File

@ -91,16 +91,6 @@ class RegistryEntryHider(StrEnum):
USER = "user"
def _convert_to_entity_category(
value: EntityCategory | str | None, raise_report: bool = True
) -> EntityCategory | None:
"""Force incoming entity_category to be an enum."""
# pylint: disable=import-outside-toplevel
from .entity import convert_to_entity_category
return convert_to_entity_category(value, raise_report=raise_report)
@attr.s(slots=True, frozen=True)
class RegistryEntry:
"""Entity Registry Entry."""
@ -115,9 +105,7 @@ class RegistryEntry:
device_id: str | None = attr.ib(default=None)
domain: str = attr.ib(init=False, repr=False)
disabled_by: RegistryEntryDisabler | None = attr.ib(default=None)
entity_category: EntityCategory | None = attr.ib(
default=None, converter=_convert_to_entity_category
)
entity_category: EntityCategory | None = attr.ib(default=None)
hidden_by: RegistryEntryHider | None = attr.ib(default=None)
icon: str | None = attr.ib(default=None)
id: str = attr.ib(factory=uuid_util.random_uuid_hex)
@ -339,8 +327,7 @@ class EntityRegistry:
capabilities: Mapping[str, Any] | None = None,
config_entry: ConfigEntry | None = None,
device_id: str | None = None,
# Type str (ENTITY_CATEG*) is deprecated as of 2021.12, use EntityCategory
entity_category: EntityCategory | str | None = None,
entity_category: EntityCategory | None = None,
original_device_class: str | None = None,
original_icon: str | None = None,
original_name: str | None = None,
@ -390,13 +377,18 @@ class EntityRegistry:
):
disabled_by = RegistryEntryDisabler.INTEGRATION
from .entity import EntityCategory # pylint: disable=import-outside-toplevel
if entity_category and not isinstance(entity_category, EntityCategory):
raise ValueError("entity_category must be a valid EntityCategory instance")
entry = RegistryEntry(
area_id=area_id,
capabilities=capabilities,
config_entry_id=config_entry_id,
device_id=device_id,
disabled_by=disabled_by,
entity_category=_convert_to_entity_category(entity_category),
entity_category=entity_category,
entity_id=entity_id,
hidden_by=hidden_by,
original_device_class=original_device_class,
@ -502,8 +494,7 @@ class EntityRegistry:
device_class: str | None | UndefinedType = UNDEFINED,
device_id: str | None | UndefinedType = UNDEFINED,
disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
# Type str (ENTITY_CATEG*) is deprecated as of 2021.12, use EntityCategory
entity_category: EntityCategory | str | None | UndefinedType = UNDEFINED,
entity_category: EntityCategory | None | UndefinedType = UNDEFINED,
hidden_by: RegistryEntryHider | None | UndefinedType = UNDEFINED,
icon: str | None | UndefinedType = UNDEFINED,
name: str | None | UndefinedType = UNDEFINED,
@ -528,6 +519,15 @@ class EntityRegistry:
):
raise ValueError("disabled_by must be a RegistryEntryDisabler value")
from .entity import EntityCategory # pylint: disable=import-outside-toplevel
if (
entity_category
and entity_category is not UNDEFINED
and not isinstance(entity_category, EntityCategory)
):
raise ValueError("entity_category must be a valid EntityCategory instance")
for attr_name, value in (
("area_id", area_id),
("capabilities", capabilities),
@ -629,6 +629,8 @@ class EntityRegistry:
)
entities = EntityRegistryItems()
from .entity import EntityCategory # pylint: disable=import-outside-toplevel
if data is not None:
for entity in data["entities"]:
# Some old installations can have some bad entities.
@ -646,9 +648,9 @@ class EntityRegistry:
disabled_by=RegistryEntryDisabler(entity["disabled_by"])
if entity["disabled_by"]
else None,
entity_category=_convert_to_entity_category(
entity["entity_category"], raise_report=False
),
entity_category=EntityCategory(entity["entity_category"])
if entity["entity_category"]
else None,
entity_id=entity["entity_id"],
hidden_by=entity["hidden_by"],
icon=entity["icon"],

View File

@ -8,6 +8,7 @@ from homeassistant.components.alexa import errors
from homeassistant.components.cloud import ALEXA_SCHEMA, alexa_config
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.entity import EntityCategory
from homeassistant.util.dt import utcnow
from tests.common import async_fire_time_changed, mock_registry
@ -28,21 +29,21 @@ async def test_alexa_config_expose_entity_prefs(hass, cloud_prefs, cloud_stub):
"test",
"light_config_id",
suggested_object_id="config_light",
entity_category="config",
entity_category=EntityCategory.CONFIG,
)
entity_entry2 = entity_registry.async_get_or_create(
"light",
"test",
"light_diagnostic_id",
suggested_object_id="diagnostic_light",
entity_category="diagnostic",
entity_category=EntityCategory.DIAGNOSTIC,
)
entity_entry3 = entity_registry.async_get_or_create(
"light",
"test",
"light_system_id",
suggested_object_id="system_light",
entity_category="system",
entity_category=EntityCategory.SYSTEM,
)
entity_entry4 = entity_registry.async_get_or_create(
"light",

View File

@ -10,6 +10,7 @@ from homeassistant.components.google_assistant import helpers as ga_helpers
from homeassistant.const import EVENT_HOMEASSISTANT_STARTED
from homeassistant.core import CoreState, State
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity import EntityCategory
from homeassistant.util.dt import utcnow
from tests.common import async_fire_time_changed, mock_registry
@ -227,21 +228,21 @@ async def test_google_config_expose_entity_prefs(hass, mock_conf, cloud_prefs):
"test",
"light_config_id",
suggested_object_id="config_light",
entity_category="config",
entity_category=EntityCategory.CONFIG,
)
entity_entry2 = entity_registry.async_get_or_create(
"light",
"test",
"light_diagnostic_id",
suggested_object_id="diagnostic_light",
entity_category="diagnostic",
entity_category=EntityCategory.DIAGNOSTIC,
)
entity_entry3 = entity_registry.async_get_or_create(
"light",
"test",
"light_system_id",
suggested_object_id="system_light",
entity_category="system",
entity_category=EntityCategory.SYSTEM,
)
entity_entry4 = entity_registry.async_get_or_create(
"light",

View File

@ -22,6 +22,7 @@ from homeassistant.components.climate import const as climate
from homeassistant.components.humidifier import const as humidifier
from homeassistant.const import CLOUD_NEVER_EXPOSED_ENTITIES
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity import EntityCategory
from . import DEMO_DEVICES
@ -136,21 +137,21 @@ async def test_sync_request(hass_fixture, assistant_client, auth_header):
"test",
"switch_config_id",
suggested_object_id="config_switch",
entity_category="config",
entity_category=EntityCategory.CONFIG,
)
entity_entry2 = entity_registry.async_get_or_create(
"switch",
"test",
"switch_diagnostic_id",
suggested_object_id="diagnostic_switch",
entity_category="diagnostic",
entity_category=EntityCategory.DIAGNOSTIC,
)
entity_entry3 = entity_registry.async_get_or_create(
"switch",
"test",
"switch_system_id",
suggested_object_id="system_switch",
entity_category="system",
entity_category=EntityCategory.SYSTEM,
)
entity_entry4 = entity_registry.async_get_or_create(
"switch",

View File

@ -14,7 +14,11 @@ from homeassistant.helpers import (
entity_platform,
entity_registry as er,
)
from homeassistant.helpers.entity import DeviceInfo, async_generate_entity_id
from homeassistant.helpers.entity import (
DeviceInfo,
EntityCategory,
async_generate_entity_id,
)
from homeassistant.helpers.entity_component import (
DEFAULT_SCAN_INTERVAL,
EntityComponent,
@ -1151,7 +1155,7 @@ async def test_entity_info_added_to_entity_registry(hass):
entity_default = MockEntity(
capability_attributes={"max": 100},
device_class="mock-device-class",
entity_category="config",
entity_category=EntityCategory.CONFIG,
icon="nice:icon",
name="best name",
supported_features=5,
@ -1170,7 +1174,7 @@ async def test_entity_info_added_to_entity_registry(hass):
"test_domain",
capabilities={"max": 100},
device_class=None,
entity_category="config",
entity_category=EntityCategory.CONFIG,
icon=None,
id=ANY,
name=None,

View File

@ -1233,26 +1233,17 @@ async def test_disabled_by_str_not_allowed(hass):
)
async def test_deprecated_entity_category_str(hass, registry, caplog):
"""Test deprecated str use of entity_category converts to enum and logs a warning."""
entry = er.RegistryEntry(
entity_id="light.kitchen",
unique_id="5678",
platform="hue",
entity_category="diagnostic",
)
async def test_entity_category_str_not_allowed(hass):
"""Test we need to pass entity category type."""
reg = er.async_get(hass)
assert entry.entity_category is EntityCategory.DIAGNOSTIC
assert " should be updated to use EntityCategory" in caplog.text
with pytest.raises(ValueError):
reg.async_get_or_create(
"light", "hue", "1234", entity_category=EntityCategory.DIAGNOSTIC.value
)
async def test_invalid_entity_category_str(hass, registry, caplog):
"""Test use of invalid entity category."""
entry = er.RegistryEntry(
entity_id="light.kitchen",
unique_id="5678",
platform="hue",
entity_category="invalid",
)
assert entry.entity_category is None
entity_id = reg.async_get_or_create("light", "hue", "1234").entity_id
with pytest.raises(ValueError):
reg.async_update_entity(
entity_id, entity_category=EntityCategory.DIAGNOSTIC.value
)

View File

@ -25,6 +25,7 @@ from homeassistant.helpers import (
template,
)
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity import EntityCategory
from homeassistant.setup import async_setup_component
from tests.common import (
@ -119,7 +120,7 @@ def area_mock(hass):
unique_id="config-in-own-area-id",
platform="test",
area_id="own-area",
entity_category="config",
entity_category=EntityCategory.CONFIG,
)
hidden_entity_in_own_area = ent_reg.RegistryEntry(
entity_id="light.hidden_in_own_area",
@ -139,7 +140,7 @@ def area_mock(hass):
unique_id="config-in-area-id",
platform="test",
device_id=device_in_area.id,
entity_category="config",
entity_category=EntityCategory.CONFIG,
)
hidden_entity_in_area = ent_reg.RegistryEntry(
entity_id="light.hidden_in_area",
@ -173,7 +174,7 @@ def area_mock(hass):
unique_id="config-no-area-id",
platform="test",
device_id=device_no_area.id,
entity_category="config",
entity_category=EntityCategory.CONFIG,
)
hidden_entity_no_area = ent_reg.RegistryEntry(
entity_id="light.hidden_no_area",