Cache entity properties that are never expected to change in the base class (#95315)

homeassistant/backports/functools.py

@@ -5,18 +5,18 @@ from collections.abc import Callable
 from types import GenericAlias
 from typing import Any, Generic, Self, TypeVar, overload
 
-_T = TypeVar("_T")
+_T_co = TypeVar("_T_co", covariant=True)
 
 
-class cached_property(Generic[_T]):
+class cached_property(Generic[_T_co]):  # pylint: disable=invalid-name
     """Backport of Python 3.12's cached_property.
 
     Includes https://github.com/python/cpython/pull/101890/files
     """
 
-    def __init__(self, func: Callable[[Any], _T]) -> None:
+    def __init__(self, func: Callable[[Any], _T_co]) -> None:
         """Initialize."""
-        self.func: Callable[[Any], _T] = func
+        self.func: Callable[[Any], _T_co] = func
         self.attrname: str | None = None
         self.__doc__ = func.__doc__
 
@@ -35,12 +35,12 @@ class cached_property(Generic[_T]):
         ...
 
     @overload
-    def __get__(self, instance: Any, owner: type[Any] | None = None) -> _T:
+    def __get__(self, instance: Any, owner: type[Any] | None = None) -> _T_co:
         ...
 
     def __get__(
         self, instance: Any | None, owner: type[Any] | None = None
-    ) -> _T | Self:
+    ) -> _T_co | Self:
         """Get."""
         if instance is None:
             return self

homeassistant/components/abode/binary_sensor.py

@@ -50,7 +50,9 @@ class AbodeBinarySensor(AbodeDevice, BinarySensorEntity):
         """Return True if the binary sensor is on."""
         return cast(bool, self._device.is_on)
 
-    @property
+    @property  # type: ignore[override]
+    # We don't know if the class may be set late here
+    # so we need to override the property to disable the cache.
     def device_class(self) -> BinarySensorDeviceClass | None:
         """Return the class of the binary sensor."""
         if self._device.get_value("is_window") == "1":

homeassistant/components/binary_sensor/__init__.py

@@ -9,6 +9,7 @@ from typing import Literal, final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
@@ -197,7 +198,7 @@ class BinarySensorEntity(Entity):
         """
         return self.device_class is not None
 
-    @property
+    @cached_property
     def device_class(self) -> BinarySensorDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/button/__init__.py

@@ -9,6 +9,7 @@ from typing import final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.config_validation import (  # noqa: F401
@@ -96,7 +97,7 @@ class ButtonEntity(RestoreEntity):
         """
         return self.device_class is not None
 
-    @property
+    @cached_property
     def device_class(self) -> ButtonDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/cover/__init__.py

@@ -11,6 +11,7 @@ from typing import Any, ParamSpec, TypeVar, final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     SERVICE_CLOSE_COVER,
@@ -250,7 +251,7 @@ class CoverEntity(Entity):
         """
         return self._attr_current_cover_tilt_position
 
-    @property
+    @cached_property
     def device_class(self) -> CoverDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/date/__init__.py

@@ -8,6 +8,7 @@ from typing import final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_DATE
 from homeassistant.core import HomeAssistant, ServiceCall
@@ -75,7 +76,7 @@ class DateEntity(Entity):
     _attr_native_value: date | None
     _attr_state: None = None
 
-    @property
+    @cached_property
     @final
     def device_class(self) -> None:
         """Return the device class for the entity."""

homeassistant/components/datetime/__init__.py

@@ -8,6 +8,7 @@ from typing import final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import config_validation as cv
@@ -86,7 +87,7 @@ class DateTimeEntity(Entity):
     _attr_state: None = None
     _attr_native_value: datetime | None
 
-    @property
+    @cached_property
     @final
     def device_class(self) -> None:
         """Return entity device class."""

homeassistant/components/dsmr/sensor.py

@@ -592,7 +592,10 @@ class DSMREntity(SensorEntity):
         """Entity is only available if there is a telegram."""
         return self.telegram is not None
 
-    @property
+    @property  # type: ignore[override]
+    # The device class can change at runtime from GAS to ENERGY
+    # when new data is received. This should be remembered and restored
+    # at startup, but the integration currently doesn't support that.
     def device_class(self) -> SensorDeviceClass | None:
         """Return the device class of this entity."""
         device_class = super().device_class

homeassistant/components/event/__init__.py

@@ -7,6 +7,7 @@ from enum import StrEnum
 import logging
 from typing import Any, Self, final
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.config_validation import (  # noqa: F401
@@ -114,7 +115,7 @@ class EventEntity(RestoreEntity):
     __last_event_type: str | None = None
     __last_event_attributes: dict[str, Any] | None = None
 
-    @property
+    @cached_property
     def device_class(self) -> EventDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/filter/sensor.py

@@ -220,10 +220,17 @@ class SensorFilter(SensorEntity):
         self._state: StateType = None
         self._filters = filters
         self._attr_icon = None
-        self._attr_device_class = None
+        self._device_class = None
         self._attr_state_class = None
         self._attr_extra_state_attributes = {ATTR_ENTITY_ID: entity_id}
 
+    @property
+    # This property is not cached because the underlying source may
+    # not always be available.
+    def device_class(self) -> SensorDeviceClass | None:  # type: ignore[override]
+        """Return the device class of the sensor."""
+        return self._device_class
+
     @callback
     def _update_filter_sensor_state_event(
         self, event: EventType[EventStateChangedData]
@@ -283,7 +290,7 @@ class SensorFilter(SensorEntity):
         self._state = temp_state.state
 
         self._attr_icon = new_state.attributes.get(ATTR_ICON, ICON)
-        self._attr_device_class = new_state.attributes.get(ATTR_DEVICE_CLASS)
+        self._device_class = new_state.attributes.get(ATTR_DEVICE_CLASS)
         self._attr_state_class = new_state.attributes.get(ATTR_STATE_CLASS)
 
         if self._attr_native_unit_of_measurement != new_state.attributes.get(

homeassistant/components/group/binary_sensor.py

@@ -5,6 +5,7 @@ from typing import Any
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.components.binary_sensor import (
     DEVICE_CLASSES_SCHEMA,
     DOMAIN as BINARY_SENSOR_DOMAIN,
@@ -147,7 +148,7 @@ class BinarySensorGroup(GroupEntity, BinarySensorEntity):
             # Set as ON if any / all member is ON
             self._attr_is_on = self.mode(state == STATE_ON for state in states)
 
-    @property
+    @cached_property
     def device_class(self) -> BinarySensorDeviceClass | None:
         """Return the sensor class of the binary sensor."""
         return self._device_class

homeassistant/components/group/sensor.py

@@ -360,7 +360,10 @@ class SensorGroup(GroupEntity, SensorEntity):
         """Return the state attributes of the sensor."""
         return {ATTR_ENTITY_ID: self._entity_ids, **self._extra_state_attribute}
 
-    @property
+    @property  # type: ignore[override]
+    # Because the device class is calculated, there is no guarantee that the
+    # sensors will be available when the entity is created so we do not want to
+    # cache the value.
     def device_class(self) -> SensorDeviceClass | None:
         """Return device class."""
         if self._attr_device_class is not None:

homeassistant/components/here_travel_time/sensor.py

@@ -154,7 +154,10 @@ class HERETravelTimeSensor(
             )
             self.async_write_ha_state()
 
-    @property
+    @property  # type: ignore[override]
+    # This property is not cached because the attribute can change
+    # at run time. This is not expected, but it is currently how
+    # the HERE integration works.
     def attribution(self) -> str | None:
         """Return the attribution."""
         if self.coordinator.data is not None:

homeassistant/components/huawei_lte/sensor.py

@@ -760,7 +760,9 @@ class HuaweiLteSensor(HuaweiLteBaseEntityWithDevice, SensorEntity):
             return self.entity_description.icon_fn(self.state)
         return self.entity_description.icon
 
-    @property
+    @property  # type: ignore[override]
+    # The device class might change at run time of the signal
+    # is not a number, so we override here.
     def device_class(self) -> SensorDeviceClass | None:
         """Return device class for sensor."""
         if self.entity_description.device_class_fn:

homeassistant/components/humidifier/__init__.py

@@ -9,6 +9,7 @@ from typing import Any, final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_MODE,
@@ -158,7 +159,7 @@ class HumidifierEntity(ToggleEntity):
 
         return data
 
-    @property
+    @cached_property
     def device_class(self) -> HumidifierDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/image_processing/__init__.py

@@ -10,6 +10,7 @@ from typing import Any, Final, TypedDict, final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.components.camera import Image
 from homeassistant.const import (
     ATTR_ENTITY_ID,
@@ -156,7 +157,7 @@ class ImageProcessingEntity(Entity):
             return self.entity_description.confidence
         return None
 
-    @property
+    @cached_property
     def device_class(self) -> ImageProcessingDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/integration/sensor.py

@@ -242,6 +242,14 @@ class IntegrationSensor(RestoreSensor):
         self._source_entity: str = source_entity
         self._last_valid_state: Decimal | None = None
         self._attr_device_info = device_info
+        self._device_class: SensorDeviceClass | None = None
+
+    @property  # type: ignore[override]
+    # The underlying source data may be unavailable at startup, so the device
+    # class may be set late so we need to override the property to disable the cache.
+    def device_class(self) -> SensorDeviceClass | None:
+        """Return the device class of the sensor."""
+        return self._device_class
 
     def _unit(self, source_unit: str) -> str:
         """Derive unit from the source sensor, SI prefix and time unit."""
@@ -288,7 +296,7 @@ class IntegrationSensor(RestoreSensor):
                         err,
                     )
 
-            self._attr_device_class = state.attributes.get(ATTR_DEVICE_CLASS)
+            self._device_class = state.attributes.get(ATTR_DEVICE_CLASS)
             self._unit_of_measurement = state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
 
         @callback
@@ -319,7 +327,7 @@ class IntegrationSensor(RestoreSensor):
                 and new_state.attributes.get(ATTR_DEVICE_CLASS)
                 == SensorDeviceClass.POWER
             ):
-                self._attr_device_class = SensorDeviceClass.ENERGY
+                self._device_class = SensorDeviceClass.ENERGY
                 self._attr_icon = None
 
             self.async_write_ha_state()

homeassistant/components/media_player/__init__.py

@@ -22,6 +22,7 @@ from aiohttp.typedefs import LooseHeaders
 import voluptuous as vol
 from yarl import URL
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.components import websocket_api
 from homeassistant.components.http import KEY_AUTHENTICATED, HomeAssistantView
 from homeassistant.components.websocket_api import ERR_NOT_SUPPORTED, ERR_UNKNOWN_ERROR
@@ -495,7 +496,7 @@ class MediaPlayerEntity(Entity):
     _attr_volume_level: float | None = None
 
     # Implement these for your media player
-    @property
+    @cached_property
     def device_class(self) -> MediaPlayerDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/mobile_app/binary_sensor.py

@@ -67,7 +67,7 @@ async def async_setup_entry(
     )
 
 
-class MobileAppBinarySensor(MobileAppEntity, BinarySensorEntity):
+class MobileAppBinarySensor(MobileAppEntity, BinarySensorEntity):  # type: ignore[misc]
     """Representation of an mobile app binary sensor."""
 
     @property

homeassistant/components/mobile_app/entity.py

@@ -69,7 +69,9 @@ class MobileAppEntity(RestoreEntity):
         """Return if entity should be enabled by default."""
         return not self._config.get(ATTR_SENSOR_DISABLED)
 
-    @property
+    @property  # type: ignore[override,unused-ignore]
+    # Because the device class is received later from the mobile app
+    # we do not want to cache the property
     def device_class(self):
         """Return the device class."""
         return self._config.get(ATTR_SENSOR_DEVICE_CLASS)

homeassistant/components/mobile_app/sensor.py

@@ -76,7 +76,7 @@ async def async_setup_entry(
     )
 
 
-class MobileAppSensor(MobileAppEntity, RestoreSensor):
+class MobileAppSensor(MobileAppEntity, RestoreSensor):  # type: ignore[misc]
     """Representation of an mobile app sensor."""
 
     async def async_restore_last_state(self, last_state):

homeassistant/components/number/__init__.py

@@ -12,6 +12,7 @@ from typing import Any, Self, final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_MODE, CONF_UNIT_OF_MEASUREMENT, UnitOfTemperature
 from homeassistant.core import HomeAssistant, ServiceCall, callback
@@ -231,7 +232,7 @@ class NumberEntity(Entity):
         """
         return self.device_class is not None
 
-    @property
+    @cached_property
     def device_class(self) -> NumberDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/sensor/__init__.py

@@ -11,6 +11,7 @@ import logging
 from math import ceil, floor, isfinite, log10
 from typing import Any, Final, Self, cast, final
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 
 # pylint: disable-next=hass-deprecated-import
@@ -259,7 +260,7 @@ class SensorEntity(Entity):
         """
         return self.device_class not in (None, SensorDeviceClass.ENUM)
 
-    @property
+    @cached_property
     def device_class(self) -> SensorDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/statistics/sensor.py

@@ -393,7 +393,9 @@ class StatisticsSensor(SensorEntity):
             unit = base_unit + "/s"
         return unit
 
-    @property
+    @property  # type: ignore[override]
+    # Since the underlying data source may not be available at startup
+    # we disable the caching of device_class.
     def device_class(self) -> SensorDeviceClass | None:
         """Return the class of this device."""
         if self._state_characteristic in STATS_DATETIME:

homeassistant/components/switch/__init__.py

@@ -8,6 +8,7 @@ import logging
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     SERVICE_TOGGLE,
@@ -102,7 +103,7 @@ class SwitchEntity(ToggleEntity):
     entity_description: SwitchEntityDescription
     _attr_device_class: SwitchDeviceClass | None
 
-    @property
+    @cached_property
     def device_class(self) -> SwitchDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/template/weather.py

@@ -294,7 +294,9 @@ class WeatherTemplate(TemplateEntity, WeatherEntity):
         """Return the daily forecast in native units."""
         return self._forecast_twice_daily
 
-    @property
+    @property  # type: ignore[override]
+    # Because attribution is a template, it can change at any time
+    # and we don't want to cache it.
     def attribution(self) -> str | None:
         """Return the attribution."""
         if self._attribution is None:

homeassistant/components/time/__init__.py

@@ -8,6 +8,7 @@ from typing import final
 
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TIME
 from homeassistant.core import HomeAssistant, ServiceCall
@@ -75,7 +76,7 @@ class TimeEntity(Entity):
     _attr_device_class: None = None
     _attr_state: None = None
 
-    @property
+    @cached_property
     @final
     def device_class(self) -> None:
         """Return the device class for the entity."""

homeassistant/components/unifiprotect/binary_sensor.py

@@ -552,6 +552,7 @@ class ProtectDeviceBinarySensor(ProtectDeviceEntity, BinarySensorEntity):
 
     device: Camera | Light | Sensor
     entity_description: ProtectBinaryEntityDescription
+    _device_class: BinarySensorDeviceClass | None
 
     @callback
     def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
@@ -561,9 +562,17 @@ class ProtectDeviceBinarySensor(ProtectDeviceEntity, BinarySensorEntity):
         self._attr_is_on = entity_description.get_ufp_value(updated_device)
         # UP Sense can be any of the 3 contact sensor device classes
         if entity_description.key == _KEY_DOOR and isinstance(updated_device, Sensor):
-            entity_description.device_class = MOUNT_DEVICE_CLASS_MAP.get(
+            self._device_class = MOUNT_DEVICE_CLASS_MAP.get(
-                updated_device.mount_type, BinarySensorDeviceClass.DOOR
+                self.device.mount_type, BinarySensorDeviceClass.DOOR
             )
+        else:
+            self._device_class = self.entity_description.device_class
+
+    @property  # type: ignore[override]
+    # UFP smart sensors can change device class at runtime
+    def device_class(self) -> BinarySensorDeviceClass | None:
+        """Return the class of this sensor."""
+        return self._device_class
 
 
 class ProtectDiskBinarySensor(ProtectNVREntity, BinarySensorEntity):

homeassistant/components/update/__init__.py

@@ -11,6 +11,7 @@ from typing import Any, Final, final
 from awesomeversion import AwesomeVersion, AwesomeVersionCompareException
 import voluptuous as vol
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.components import websocket_api
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_OFF, STATE_ON, EntityCategory
@@ -223,7 +224,7 @@ class UpdateEntity(RestoreEntity):
         """
         return self.device_class is not None
 
-    @property
+    @cached_property
     def device_class(self) -> UpdateDeviceClass | None:
         """Return the class of this entity."""
         if hasattr(self, "_attr_device_class"):

homeassistant/components/zha/binary_sensor.py

@@ -8,6 +8,7 @@ import zigpy.types as t
 from zigpy.zcl.clusters.general import OnOff
 from zigpy.zcl.clusters.security import IasZone
 
+from homeassistant.backports.functools import cached_property
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
@@ -195,7 +196,7 @@ class IASZone(BinarySensor):
         zone_type = self._cluster_handler.cluster.get("zone_type")
         return IAS_ZONE_NAME_MAPPING.get(zone_type, "iaszone")
 
-    @property
+    @cached_property
     def device_class(self) -> BinarySensorDeviceClass | None:
         """Return device class from component DEVICE_CLASSES."""
         zone_type = self._cluster_handler.cluster.get("zone_type")

homeassistant/components/zwave_js/sensor.py

@@ -645,6 +645,13 @@ class ZwaveSensor(ZWaveBaseEntity, SensorEntity):
             return None
         return str(self.info.primary_value.metadata.unit)
 
+    @property  # type: ignore[override]
+    # fget is used in the child classes which is not compatible with cached_property
+    # mypy also doesn't know about fget: https://github.com/python/mypy/issues/6185
+    def device_class(self) -> SensorDeviceClass | None:
+        """Return device class of sensor."""
+        return super().device_class
+
 
 class ZWaveNumericSensor(ZwaveSensor):
     """Representation of a Z-Wave Numeric sensor."""
@@ -737,7 +744,9 @@ class ZWaveListSensor(ZwaveSensor):
             return list(self.info.primary_value.metadata.states.values())
         return None
 
-    @property
+    @property  # type: ignore[override]
+    # fget is used which is not compatible with cached_property
+    # mypy also doesn't know about fget: https://github.com/python/mypy/issues/6185
     def device_class(self) -> SensorDeviceClass | None:
         """Return sensor device class."""
         if (device_class := super().device_class) is not None:
@@ -781,7 +790,7 @@ class ZWaveConfigParameterSensor(ZWaveListSensor):
             additional_info=[property_key_name] if property_key_name else None,
         )
 
-    @property
+    @property  # type: ignore[override]
     def device_class(self) -> SensorDeviceClass | None:
         """Return sensor device class."""
         # mypy doesn't know about fget: https://github.com/python/mypy/issues/6185

homeassistant/helpers/entity.py

@@ -550,7 +550,7 @@ class Entity(ABC):
         """
         return self._attr_device_info
 
-    @property
+    @cached_property
     def device_class(self) -> str | None:
         """Return the class of this device, from component DEVICE_CLASSES."""
         if hasattr(self, "_attr_device_class"):
@@ -639,7 +639,7 @@ class Entity(ABC):
             return self.entity_description.entity_registry_visible_default
         return True
 
-    @property
+    @cached_property
     def attribution(self) -> str | None:
         """Return the attribution."""
         return self._attr_attribution
@@ -653,7 +653,7 @@ class Entity(ABC):
             return self.entity_description.entity_category
         return None
 
-    @property
+    @cached_property
     def translation_key(self) -> str | None:
         """Return the translation key to translate the entity's states."""
         if hasattr(self, "_attr_translation_key"):

tests/components/event/test_init.py

@@ -51,6 +51,7 @@ async def test_event() -> None:
         event.event_types
 
     # Test retrieving data from entity description
+    del event.device_class
     event.entity_description = EventEntityDescription(
         key="test_event",
         event_types=["short_press", "long_press"],
@@ -63,6 +64,7 @@ async def test_event() -> None:
     event._attr_event_types = ["short_press", "long_press", "double_press"]
     assert event.event_types == ["short_press", "long_press", "double_press"]
     event._attr_device_class = EventDeviceClass.BUTTON
+    del event.device_class
     assert event.device_class == EventDeviceClass.BUTTON
 
     # Test triggering an event

tests/components/update/test_init.py

@@ -59,11 +59,13 @@ class MockUpdateEntity(UpdateEntity):
     """Mock UpdateEntity to use in tests."""
 
 
-async def test_update(hass: HomeAssistant) -> None:
+def _create_mock_update_entity(
-    """Test getting data from the mocked update entity."""
+    hass: HomeAssistant,
+) -> MockUpdateEntity:
+    mock_platform = MockEntityPlatform(hass)
     update = MockUpdateEntity()
     update.hass = hass
-    update.platform = MockEntityPlatform(hass)
+    update.platform = mock_platform
 
     update._attr_installed_version = "1.0.0"
     update._attr_latest_version = "1.0.1"
@@ -71,6 +73,13 @@ async def test_update(hass: HomeAssistant) -> None:
     update._attr_release_url = "https://example.com"
     update._attr_title = "Title"
 
+    return update
+
+
+async def test_update(hass: HomeAssistant) -> None:
+    """Test getting data from the mocked update entity."""
+    update = _create_mock_update_entity(hass)
+
     assert update.entity_category is EntityCategory.DIAGNOSTIC
     assert (
         update.entity_picture
@@ -93,7 +102,6 @@ async def test_update(hass: HomeAssistant) -> None:
         ATTR_SKIPPED_VERSION: None,
         ATTR_TITLE: "Title",
     }
-
     # Test no update available
     update._attr_installed_version = "1.0.0"
     update._attr_latest_version = "1.0.0"
@@ -120,14 +128,19 @@ async def test_update(hass: HomeAssistant) -> None:
     assert update.state is STATE_ON
 
     # Test entity category becomes config when its possible to install
+    update = _create_mock_update_entity(hass)
     update._attr_supported_features = UpdateEntityFeature.INSTALL
     assert update.entity_category is EntityCategory.CONFIG
 
     # UpdateEntityDescription was set
+    update = _create_mock_update_entity(hass)
     update._attr_supported_features = 0
     update.entity_description = UpdateEntityDescription(key="F5 - Its very refreshing")
     assert update.device_class is None
     assert update.entity_category is EntityCategory.CONFIG
+
+    update = _create_mock_update_entity(hass)
+    update._attr_supported_features = 0
     update.entity_description = UpdateEntityDescription(
         key="F5 - Its very refreshing",
         device_class=UpdateDeviceClass.FIRMWARE,
@@ -137,14 +150,24 @@ async def test_update(hass: HomeAssistant) -> None:
     assert update.entity_category is None
 
     # Device class via attribute (override entity description)
+    update = _create_mock_update_entity(hass)
+    update._attr_supported_features = 0
     update._attr_device_class = None
     assert update.device_class is None
+
+    update = _create_mock_update_entity(hass)
+    update._attr_supported_features = 0
     update._attr_device_class = UpdateDeviceClass.FIRMWARE
     assert update.device_class is UpdateDeviceClass.FIRMWARE
 
     # Entity Attribute via attribute (override entity description)
+    update = _create_mock_update_entity(hass)
+    update._attr_supported_features = 0
     update._attr_entity_category = None
     assert update.entity_category is None
+
+    update = _create_mock_update_entity(hass)
+    update._attr_supported_features = 0
     update._attr_entity_category = EntityCategory.DIAGNOSTIC
     assert update.entity_category is EntityCategory.DIAGNOSTIC
 

tests/helpers/test_entity.py

@@ -98,9 +98,13 @@ class TestHelpersEntity:
 
     def setup_method(self, method):
         """Set up things to be run when tests are started."""
+        self.hass = get_test_home_assistant()
+        self._create_entity()
+
+    def _create_entity(self) -> None:
         self.entity = entity.Entity()
         self.entity.entity_id = "test.overwrite_hidden_true"
-        self.hass = self.entity.hass = get_test_home_assistant()
+        self.entity.hass = self.hass
         self.entity.schedule_update_ha_state()
         self.hass.block_till_done()
 
@@ -123,6 +127,7 @@ class TestHelpersEntity:
         with patch(
             "homeassistant.helpers.entity.Entity.device_class", new="test_class"
         ):
+            self._create_entity()
             self.entity.schedule_update_ha_state()
             self.hass.block_till_done()
         state = self.hass.states.get(self.entity.entity_id)