From 7fc2d9e087a906ef1121c121ffc2d69f86ddec72 Mon Sep 17 00:00:00 2001 From: Christopher Bailey Date: Sun, 14 Aug 2022 16:57:25 -0400 Subject: [PATCH] Persist previous mic/record values for UniFi Protect privacy mode (#76472) --- .../components/unifiprotect/switch.py | 140 +++++++++++++----- tests/components/unifiprotect/test_switch.py | 41 ++++- 2 files changed, 133 insertions(+), 48 deletions(-) diff --git a/homeassistant/components/unifiprotect/switch.py b/homeassistant/components/unifiprotect/switch.py index 5bc4e1f17eb..e0ddddd4c53 100644 --- a/homeassistant/components/unifiprotect/switch.py +++ b/homeassistant/components/unifiprotect/switch.py @@ -2,22 +2,23 @@ from __future__ import annotations from dataclasses import dataclass -import logging from typing import Any from pyunifiprotect.data import ( Camera, ProtectAdoptableDeviceModel, + ProtectModelWithId, RecordingMode, VideoMode, ) from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.config_entries import ConfigEntry -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.restore_state import RestoreEntity from .const import DISPATCH_ADOPT, DOMAIN from .data import ProtectData @@ -25,7 +26,8 @@ from .entity import ProtectDeviceEntity, async_all_device_entities from .models import PermRequired, ProtectSetableKeysMixin, T from .utils import async_dispatch_id as _ufpd -_LOGGER = logging.getLogger(__name__) +ATTR_PREV_MIC = "prev_mic_level" +ATTR_PREV_RECORD = "prev_record_mode" @dataclass @@ -35,9 +37,6 @@ class ProtectSwitchEntityDescription( """Describes UniFi Protect Switch entity.""" -_KEY_PRIVACY_MODE = "privacy_mode" - - async def _set_highfps(obj: Camera, value: bool) -> None: if value: await obj.set_video_mode(VideoMode.HIGH_FPS) @@ -86,15 +85,6 @@ CAMERA_SWITCHES: tuple[ProtectSwitchEntityDescription, ...] = ( ufp_set_method_fn=_set_highfps, ufp_perm=PermRequired.WRITE, ), - ProtectSwitchEntityDescription( - key=_KEY_PRIVACY_MODE, - name="Privacy Mode", - icon="mdi:eye-settings", - entity_category=EntityCategory.CONFIG, - ufp_required_field="feature_flags.has_privacy_mask", - ufp_value="is_privacy_on", - ufp_perm=PermRequired.WRITE, - ), ProtectSwitchEntityDescription( key="system_sounds", name="System Sounds", @@ -192,6 +182,16 @@ CAMERA_SWITCHES: tuple[ProtectSwitchEntityDescription, ...] = ( ), ) +PRIVACY_MODE_SWITCH = ProtectSwitchEntityDescription[Camera]( + key="privacy_mode", + name="Privacy Mode", + icon="mdi:eye-settings", + entity_category=EntityCategory.CONFIG, + ufp_required_field="feature_flags.has_privacy_mask", + ufp_value="is_privacy_on", + ufp_perm=PermRequired.WRITE, +) + SENSE_SWITCHES: tuple[ProtectSwitchEntityDescription, ...] = ( ProtectSwitchEntityDescription( key="status_light", @@ -316,6 +316,11 @@ async def async_setup_entry( viewer_descs=VIEWER_SWITCHES, ufp_device=device, ) + entities += async_all_device_entities( + data, + ProtectPrivacyModeSwitch, + camera_descs=[PRIVACY_MODE_SWITCH], + ) async_add_entities(entities) entry.async_on_unload( @@ -331,6 +336,11 @@ async def async_setup_entry( lock_descs=DOORLOCK_SWITCHES, viewer_descs=VIEWER_SWITCHES, ) + entities += async_all_device_entities( + data, + ProtectPrivacyModeSwitch, + camera_descs=[PRIVACY_MODE_SWITCH], + ) async_add_entities(entities) @@ -350,17 +360,6 @@ class ProtectSwitch(ProtectDeviceEntity, SwitchEntity): self._attr_name = f"{self.device.display_name} {self.entity_description.name}" self._switch_type = self.entity_description.key - if not isinstance(self.device, Camera): - return - - if self.entity_description.key == _KEY_PRIVACY_MODE: - if self.device.is_privacy_on: - self._previous_mic_level = 100 - self._previous_record_mode = RecordingMode.ALWAYS - else: - self._previous_mic_level = self.device.mic_volume - self._previous_record_mode = self.device.recording_settings.mode - @property def is_on(self) -> bool: """Return true if device is on.""" @@ -368,24 +367,83 @@ class ProtectSwitch(ProtectDeviceEntity, SwitchEntity): async def async_turn_on(self, **kwargs: Any) -> None: """Turn the device on.""" - if self._switch_type == _KEY_PRIVACY_MODE: - assert isinstance(self.device, Camera) - self._previous_mic_level = self.device.mic_volume - self._previous_record_mode = self.device.recording_settings.mode - await self.device.set_privacy(True, 0, RecordingMode.NEVER) - else: - await self.entity_description.ufp_set(self.device, True) + + await self.entity_description.ufp_set(self.device, True) async def async_turn_off(self, **kwargs: Any) -> None: """Turn the device off.""" - if self._switch_type == _KEY_PRIVACY_MODE: - assert isinstance(self.device, Camera) - _LOGGER.debug( - "Setting Privacy Mode to false for %s", self.device.display_name - ) - await self.device.set_privacy( - False, self._previous_mic_level, self._previous_record_mode + await self.entity_description.ufp_set(self.device, False) + + +class ProtectPrivacyModeSwitch(RestoreEntity, ProtectSwitch): + """A UniFi Protect Switch.""" + + device: Camera + + def __init__( + self, + data: ProtectData, + device: ProtectAdoptableDeviceModel, + description: ProtectSwitchEntityDescription, + ) -> None: + """Initialize an UniFi Protect Switch.""" + super().__init__(data, device, description) + + if self.device.is_privacy_on: + extra_state = self.extra_state_attributes or {} + self._previous_mic_level = extra_state.get(ATTR_PREV_MIC, 100) + self._previous_record_mode = extra_state.get( + ATTR_PREV_RECORD, RecordingMode.ALWAYS ) else: - await self.entity_description.ufp_set(self.device, False) + self._previous_mic_level = self.device.mic_volume + self._previous_record_mode = self.device.recording_settings.mode + + @callback + def _update_previous_attr(self) -> None: + if self.is_on: + self._attr_extra_state_attributes = { + ATTR_PREV_MIC: self._previous_mic_level, + ATTR_PREV_RECORD: self._previous_record_mode, + } + else: + self._attr_extra_state_attributes = {} + + @callback + def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None: + super()._async_update_device_from_protect(device) + + # do not add extra state attribute on initialize + if self.entity_id: + self._update_previous_attr() + + async def async_turn_on(self, **kwargs: Any) -> None: + """Turn the device on.""" + + self._previous_mic_level = self.device.mic_volume + self._previous_record_mode = self.device.recording_settings.mode + await self.device.set_privacy(True, 0, RecordingMode.NEVER) + + async def async_turn_off(self, **kwargs: Any) -> None: + """Turn the device off.""" + + extra_state = self.extra_state_attributes or {} + prev_mic = extra_state.get(ATTR_PREV_MIC, self._previous_mic_level) + prev_record = extra_state.get(ATTR_PREV_RECORD, self._previous_record_mode) + await self.device.set_privacy(False, prev_mic, prev_record) + + async def async_added_to_hass(self) -> None: + """Restore extra state attributes on startp up.""" + await super().async_added_to_hass() + + if not (last_state := await self.async_get_last_state()): + return + + self._previous_mic_level = last_state.attributes.get( + ATTR_PREV_MIC, self._previous_mic_level + ) + self._previous_record_mode = last_state.attributes.get( + ATTR_PREV_RECORD, self._previous_record_mode + ) + self._update_previous_attr() diff --git a/tests/components/unifiprotect/test_switch.py b/tests/components/unifiprotect/test_switch.py index 684e3b8e441..6fa718c4952 100644 --- a/tests/components/unifiprotect/test_switch.py +++ b/tests/components/unifiprotect/test_switch.py @@ -9,8 +9,11 @@ from pyunifiprotect.data import Camera, Light, Permission, RecordingMode, VideoM from homeassistant.components.unifiprotect.const import DEFAULT_ATTRIBUTION from homeassistant.components.unifiprotect.switch import ( + ATTR_PREV_MIC, + ATTR_PREV_RECORD, CAMERA_SWITCHES, LIGHT_SWITCHES, + PRIVACY_MODE_SWITCH, ProtectSwitchEntityDescription, ) from homeassistant.const import ATTR_ATTRIBUTION, ATTR_ENTITY_ID, STATE_OFF, Platform @@ -347,31 +350,55 @@ async def test_switch_camera_highfps( async def test_switch_camera_privacy( hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera ): - """Tests Privacy Mode switch for cameras.""" + """Tests Privacy Mode switch for cameras with privacy mode defaulted on.""" + + previous_mic = doorbell.mic_volume = 53 + previous_record = doorbell.recording_settings.mode = RecordingMode.DETECTIONS await init_entry(hass, ufp, [doorbell]) assert_entity_counts(hass, Platform.SWITCH, 13, 12) - description = CAMERA_SWITCHES[4] + description = PRIVACY_MODE_SWITCH doorbell.__fields__["set_privacy"] = Mock() doorbell.set_privacy = AsyncMock() _, entity_id = ids_from_device_description(Platform.SWITCH, doorbell, description) + state = hass.states.get(entity_id) + assert state and state.state == "off" + assert ATTR_PREV_MIC not in state.attributes + assert ATTR_PREV_RECORD not in state.attributes + await hass.services.async_call( "switch", "turn_on", {ATTR_ENTITY_ID: entity_id}, blocking=True ) - doorbell.set_privacy.assert_called_once_with(True, 0, RecordingMode.NEVER) + doorbell.set_privacy.assert_called_with(True, 0, RecordingMode.NEVER) + + new_doorbell = doorbell.copy() + new_doorbell.add_privacy_zone() + new_doorbell.mic_volume = 0 + new_doorbell.recording_settings.mode = RecordingMode.NEVER + ufp.api.bootstrap.cameras = {new_doorbell.id: new_doorbell} + + mock_msg = Mock() + mock_msg.changed_data = {} + mock_msg.new_obj = new_doorbell + ufp.ws_msg(mock_msg) + + state = hass.states.get(entity_id) + assert state and state.state == "on" + assert state.attributes[ATTR_PREV_MIC] == previous_mic + assert state.attributes[ATTR_PREV_RECORD] == previous_record.value + + doorbell.set_privacy.reset_mock() await hass.services.async_call( "switch", "turn_off", {ATTR_ENTITY_ID: entity_id}, blocking=True ) - doorbell.set_privacy.assert_called_with( - False, doorbell.mic_volume, doorbell.recording_settings.mode - ) + doorbell.set_privacy.assert_called_with(False, previous_mic, previous_record) async def test_switch_camera_privacy_already_on( @@ -383,7 +410,7 @@ async def test_switch_camera_privacy_already_on( await init_entry(hass, ufp, [doorbell]) assert_entity_counts(hass, Platform.SWITCH, 13, 12) - description = CAMERA_SWITCHES[4] + description = PRIVACY_MODE_SWITCH doorbell.__fields__["set_privacy"] = Mock() doorbell.set_privacy = AsyncMock()