From 151b3b3b0a86f148d9e609a4330c6505b5b6e732 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 16 Jun 2024 14:14:59 -0500
Subject: [PATCH] Reduce duplicate code in unifiprotect entities (#119779)

---
 .../components/unifiprotect/sensor.py         | 13 ++---
 .../components/unifiprotect/switch.py         | 53 ++++++++-----------
 2 files changed, 26 insertions(+), 40 deletions(-)

diff --git a/homeassistant/components/unifiprotect/sensor.py b/homeassistant/components/unifiprotect/sensor.py
index ea2c84bb128..78fc24026a3 100644
--- a/homeassistant/components/unifiprotect/sensor.py
+++ b/homeassistant/components/unifiprotect/sensor.py
@@ -706,8 +706,8 @@ def _async_nvr_entities(
     return entities
 
 
-class ProtectDeviceSensor(ProtectDeviceEntity, SensorEntity):
-    """A Ubiquiti UniFi Protect Sensor."""
+class BaseProtectSensor(BaseProtectEntity, SensorEntity):
+    """A UniFi Protect Sensor Entity."""
 
     entity_description: ProtectSensorEntityDescription
     _state_attrs = ("_attr_available", "_attr_native_value")
@@ -717,15 +717,12 @@ class ProtectDeviceSensor(ProtectDeviceEntity, SensorEntity):
         self._attr_native_value = self.entity_description.get_ufp_value(self.device)
 
 
-class ProtectNVRSensor(ProtectNVREntity, SensorEntity):
+class ProtectDeviceSensor(BaseProtectSensor, ProtectDeviceEntity):
     """A Ubiquiti UniFi Protect Sensor."""
 
-    entity_description: ProtectSensorEntityDescription
-    _state_attrs = ("_attr_available", "_attr_native_value")
 
-    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
-        super()._async_update_device_from_protect(device)
-        self._attr_native_value = self.entity_description.get_ufp_value(self.device)
+class ProtectNVRSensor(BaseProtectSensor, ProtectNVREntity):
+    """A Ubiquiti UniFi Protect Sensor."""
 
 
 class ProtectEventSensor(EventEntityMixin, SensorEntity):
diff --git a/homeassistant/components/unifiprotect/switch.py b/homeassistant/components/unifiprotect/switch.py
index 104c8f4af86..ca56a602209 100644
--- a/homeassistant/components/unifiprotect/switch.py
+++ b/homeassistant/components/unifiprotect/switch.py
@@ -472,43 +472,32 @@ _PRIVACY_DESCRIPTIONS: dict[ModelType, Sequence[ProtectEntityDescription]] = {
 }
 
 
-class ProtectSwitch(ProtectDeviceEntity, SwitchEntity):
+class ProtectBaseSwitch(BaseProtectEntity, SwitchEntity):
+    """Base class for UniFi Protect Switch."""
+
+    entity_description: ProtectSwitchEntityDescription
+    _state_attrs = ("_attr_available", "_attr_is_on")
+
+    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) is True
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the device on."""
+        await self.entity_description.ufp_set(self.device, True)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the device off."""
+        await self.entity_description.ufp_set(self.device, False)
+
+
+class ProtectSwitch(ProtectBaseSwitch, ProtectDeviceEntity):
     """A UniFi Protect Switch."""
 
-    entity_description: ProtectSwitchEntityDescription
-    _state_attrs = ("_attr_available", "_attr_is_on")
 
-    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) is True
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn the device on."""
-        await self.entity_description.ufp_set(self.device, True)
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn the device off."""
-        await self.entity_description.ufp_set(self.device, False)
-
-
-class ProtectNVRSwitch(ProtectNVREntity, SwitchEntity):
+class ProtectNVRSwitch(ProtectBaseSwitch, ProtectNVREntity):
     """A UniFi Protect NVR Switch."""
 
-    entity_description: ProtectSwitchEntityDescription
-    _state_attrs = ("_attr_available", "_attr_is_on")
-
-    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) is True
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn the device on."""
-        await self.entity_description.ufp_set(self.device, True)
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn the device off."""
-        await self.entity_description.ufp_set(self.device, False)
-
 
 class ProtectPrivacyModeSwitch(RestoreEntity, ProtectSwitch):
     """A UniFi Protect Switch."""