Cleanup unifiprotect entity classes

homeassistant/components/unifiprotect/binary_sensor.py

@@ -8,11 +8,9 @@
 from uiprotect.data import (
     NVR,
     Camera,
-    Light,
     ModelType,
     MountType,
     ProtectAdoptableDeviceModel,
-    ProtectModelWithId,
     Sensor,
     SmartDetectObjectType,
 )
@@ -27,11 +25,12 @@
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .data import ProtectData, UFPConfigEntry
+from .data import ProtectData, ProtectDeviceType, UFPConfigEntry
 from .entity import (
     BaseProtectEntity,
     EventEntityMixin,
     ProtectDeviceEntity,
+    ProtectIsOnEntity,
     ProtectNVREntity,
     async_all_device_entities,
 )
@@ -623,31 +622,22 @@ class ProtectBinaryEventEntityDescription(
 }
 
 
-class ProtectDeviceBinarySensor(ProtectDeviceEntity, BinarySensorEntity):
+class ProtectDeviceBinarySensor(
+    ProtectIsOnEntity, ProtectDeviceEntity, BinarySensorEntity
+):
     """A UniFi Protect Device Binary Sensor."""
 
-    device: Camera | Light | Sensor
     entity_description: ProtectBinaryEntityDescription
-    _state_attrs: tuple[str, ...] = ("_attr_available", "_attr_is_on")
-
-    @callback
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
-        super()._async_update_device_from_protect(device)
-        self._attr_is_on = self.entity_description.get_ufp_value(self.device)
 
 
 class MountableProtectDeviceBinarySensor(ProtectDeviceBinarySensor):
     """A UniFi Protect Device Binary Sensor that can change device class at runtime."""
 
     device: Sensor
-    _state_attrs: tuple[str, ...] = (
-        "_attr_available",
-        "_attr_is_on",
-        "_attr_device_class",
-    )
+    _state_attrs = ("_attr_available", "_attr_is_on", "_attr_device_class")
 
     @callback
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
+    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
         super()._async_update_device_from_protect(device)
         # UP Sense can be any of the 3 contact sensor device classes
         self._attr_device_class = MOUNT_DEVICE_CLASS_MAP.get(
@@ -673,7 +663,6 @@ def __init__(
         self._disk = disk
         # backwards compat with old unique IDs
         index = self._disk.slot - 1
-
         description = dataclasses.replace(
             description,
             key=f"{description.key}_{index}",
@@ -682,7 +671,7 @@ def __init__(
         super().__init__(data, device, description)
 
     @callback
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
+    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
         super()._async_update_device_from_protect(device)
         slot = self._disk.slot
         self._attr_available = False
@@ -712,7 +701,7 @@ def _set_event_done(self) -> None:
         self._attr_extra_state_attributes = {}
 
     @callback
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
+    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
         description = self.entity_description
 
         prev_event = self._event

homeassistant/components/unifiprotect/button.py

@@ -4,10 +4,11 @@
 
 from collections.abc import Sequence
 from dataclasses import dataclass
+from functools import partial
 import logging
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
-from uiprotect.data import ModelType, ProtectAdoptableDeviceModel, ProtectModelWithId
+from uiprotect.data import ModelType, ProtectAdoptableDeviceModel
 
 from homeassistant.components.button import (
     ButtonDeviceClass,
@@ -21,7 +22,7 @@
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DEVICES_THAT_ADOPT, DOMAIN
-from .data import UFPConfigEntry
+from .data import ProtectDeviceType, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectEntityDescription, ProtectSetableKeysMixin, T
 
@@ -38,7 +39,6 @@ class ProtectButtonEntityDescription(
 
 
 DEVICE_CLASS_CHIME_BUTTON: Final = "unifiprotect__chime_button"
-KEY_ADOPT = "adopt"
 
 
 ALL_DEVICE_BUTTONS: tuple[ProtectButtonEntityDescription, ...] = (
@@ -61,7 +61,7 @@ class ProtectButtonEntityDescription(
 )
 
 ADOPT_BUTTON = ProtectButtonEntityDescription[ProtectAdoptableDeviceModel](
-    key=KEY_ADOPT,
+    key="adopt",
     name="Adopt device",
     icon="mdi:plus-circle",
     ufp_press="adopt",
@@ -119,47 +119,39 @@ async def async_setup_entry(
     """Discover devices on a UniFi Protect NVR."""
     data = entry.runtime_data
 
+    adopt_entities = partial(
+        async_all_device_entities,
+        data,
+        ProtectAdoptButton,
+        unadopted_descs=[ADOPT_BUTTON],
+    )
+    base_entities = partial(
+        async_all_device_entities,
+        data,
+        ProtectButton,
+        all_descs=ALL_DEVICE_BUTTONS,
+        model_descriptions=_MODEL_DESCRIPTIONS,
+    )
+
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:
-        entities = async_all_device_entities(
-            data,
-            ProtectButton,
-            all_descs=ALL_DEVICE_BUTTONS,
-            unadopted_descs=[ADOPT_BUTTON],
-            model_descriptions=_MODEL_DESCRIPTIONS,
-            ufp_device=device,
+        async_add_entities(
+            [*base_entities(ufp_device=device), *adopt_entities(ufp_device=device)]
         )
-        async_add_entities(entities)
         _async_remove_adopt_button(hass, device)
 
     @callback
     def _async_add_unadopted_device(device: ProtectAdoptableDeviceModel) -> None:
         if not device.can_adopt or not device.can_create(data.api.bootstrap.auth_user):
             _LOGGER.debug("Device is not adoptable: %s", device.id)
             return
-        async_add_entities(
-            async_all_device_entities(
-                data,
-                ProtectButton,
-                unadopted_descs=[ADOPT_BUTTON],
-                ufp_device=device,
-            )
-        )
+        async_add_entities(adopt_entities(ufp_device=device))
 
     data.async_subscribe_adopt(_add_new_device)
     entry.async_on_unload(
         async_dispatcher_connect(hass, data.add_signal, _async_add_unadopted_device)
     )
-
-    async_add_entities(
-        async_all_device_entities(
-            data,
-            ProtectButton,
-            all_descs=ALL_DEVICE_BUTTONS,
-            unadopted_descs=[ADOPT_BUTTON],
-            model_descriptions=_MODEL_DESCRIPTIONS,
-        )
-    )
+    async_add_entities([*base_entities(), *adopt_entities()])
 
     for device in data.get_by_types(DEVICES_THAT_ADOPT):
         _async_remove_adopt_button(hass, device)
@@ -170,16 +162,20 @@ class ProtectButton(ProtectDeviceEntity, ButtonEntity):
 
     entity_description: ProtectButtonEntityDescription
 
-    @callback
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
-        super()._async_update_device_from_protect(device)
-        if self.entity_description.key == KEY_ADOPT:
-            device = self.device
-            self._attr_available = device.can_adopt and device.can_create(
-                self.data.api.bootstrap.auth_user
-            )
-
     async def async_press(self) -> None:
         """Press the button."""
         if self.entity_description.ufp_press is not None:
             await getattr(self.device, self.entity_description.ufp_press)()
+
+
+class ProtectAdoptButton(ProtectButton):
+    """A Ubiquiti UniFi Protect Adopt button."""
+
+    @callback
+    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
+        super()._async_update_device_from_protect(device)
+        if TYPE_CHECKING:
+            assert isinstance(device, ProtectAdoptableDeviceModel)
+        self._attr_available = device.can_adopt and device.can_create(
+            self.data.api.bootstrap.auth_user
+        )

homeassistant/components/unifiprotect/camera.py

@@ -9,7 +9,6 @@
     Camera as UFPCamera,
     CameraChannel,
     ProtectAdoptableDeviceModel,
-    ProtectModelWithId,
     StateType,
 )
 
@@ -28,7 +27,7 @@
     ATTR_WIDTH,
     DOMAIN,
 )
-from .data import ProtectData, UFPConfigEntry
+from .data import ProtectData, ProtectDeviceType, UFPConfigEntry
 from .entity import ProtectDeviceEntity
 from .utils import get_camera_base_name
 
@@ -216,7 +215,7 @@ def _async_set_stream_source(self) -> None:
             self._attr_supported_features = _EMPTY_CAMERA_FEATURES
 
     @callback
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
+    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
         super()._async_update_device_from_protect(device)
         updated_device = self.device
         channel = updated_device.channels[self.channel.id]