evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Handle remove packets for UniFi Protect (#77337)

Browse source 
Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
Christopher Bailey 2022-08-25 23:05:18 -04:00 • committed by GitHub
parent 47848b7cf8
commit 6558643448
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 52 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

5
homeassistant/components/unifiprotect/button.py
View file

@ -129,7 +129,8 @@ async def async_setup_entry(
async_add_entities(entities)
_async_remove_adopt_button(hass, device)
async def _add_unadopted_device(device: ProtectAdoptableDeviceModel) -> None:
@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
@ -147,7 +148,7 @@ async def async_setup_entry(
)
entry.async_on_unload(
async_dispatcher_connect(
hass, _ufpd(entry, DISPATCH_ADD), _add_unadopted_device
hass, _ufpd(entry, DISPATCH_ADD), _async_add_unadopted_device
)
)

16
homeassistant/components/unifiprotect/data.py
View file

@ -21,6 +21,7 @@ from pyunifiprotect.exceptions import ClientError, NotAuthorized
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.event import async_track_time_interval
@ -168,6 +169,18 @@ class ProtectData:
_LOGGER.debug("New device detected: %s", device.id)
async_dispatcher_send(self._hass, _ufpd(self._entry, DISPATCH_ADD), device)
@callback
def _async_remove_device(self, device: ProtectAdoptableDeviceModel) -> None:
registry = dr.async_get(self._hass)
device_entry = registry.async_get_device(
identifiers=set(), connections={(dr.CONNECTION_NETWORK_MAC, device.mac)}
)
if device_entry:
_LOGGER.debug("Device removed: %s", device.id)
registry.async_update_device(
device_entry.id, remove_config_entry_id=self._entry.entry_id
)
@callback
def _async_update_device(
self, device: ProtectAdoptableDeviceModel | NVR, changed_data: dict[str, Any]
@ -195,8 +208,9 @@ class ProtectData:
@callback
def _async_process_ws_message(self, message: WSSubscriptionMessage) -> None:
# removed packets are not processed yet
if message.new_obj is None:
if isinstance(message.old_obj, ProtectAdoptableDeviceModel):
self._async_remove_device(message.old_obj)
return
obj = message.new_obj

4
homeassistant/components/unifiprotect/entity.py
View file

@ -41,7 +41,7 @@ def _async_device_entities(
unadopted_descs: Sequence[ProtectRequiredKeysMixin],
ufp_device: ProtectAdoptableDeviceModel | None = None,
) -> list[ProtectDeviceEntity]:
if len(descs) + len(unadopted_descs) == 0:
if not descs and not unadopted_descs:
return []
entities: list[ProtectDeviceEntity] = []
@ -162,7 +162,7 @@ def async_all_device_entities(
elif ufp_device.model == ModelType.CHIME:
descs = chime_descs
if len(descs) + len(unadopted_descs) == 0 or ufp_device.model is None:
if not descs and not unadopted_descs or ufp_device.model is None:
return []
return _async_device_entities(
data, klass, ufp_device.model, descs, unadopted_descs, ufp_device

6
tests/components/unifiprotect/test_binary_sensor.py
View file

@ -51,7 +51,7 @@ async def test_binary_sensor_camera_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 3, 3)
await remove_entities(hass, [doorbell, unadopted_camera])
await remove_entities(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0)
await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 3, 3)
@ -65,7 +65,7 @@ async def test_binary_sensor_light_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [light])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 2, 2)
await remove_entities(hass, [light])
await remove_entities(hass, ufp, [light])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0)
await adopt_devices(hass, ufp, [light])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 2, 2)
@ -79,7 +79,7 @@ async def test_binary_sensor_sensor_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [sensor_all])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 4, 4)
await remove_entities(hass, [sensor_all])
await remove_entities(hass, ufp, [sensor_all])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0)
await adopt_devices(hass, ufp, [sensor_all])
assert_entity_counts(hass, Platform.BINARY_SENSOR, 4, 4)

2
tests/components/unifiprotect/test_button.py
View file

@ -28,7 +28,7 @@ async def test_button_chime_remove(
await init_entry(hass, ufp, [chime])
assert_entity_counts(hass, Platform.BUTTON, 4, 2)
await remove_entities(hass, [chime])
await remove_entities(hass, ufp, [chime])
assert_entity_counts(hass, Platform.BUTTON, 0, 0)
await adopt_devices(hass, ufp, [chime])
assert_entity_counts(hass, Platform.BUTTON, 4, 2)

4
tests/components/unifiprotect/test_camera.py
View file

@ -279,7 +279,7 @@ async def test_adopt(hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCa
await init_entry(hass, ufp, [camera1])
assert_entity_counts(hass, Platform.CAMERA, 0, 0)
await remove_entities(hass, [camera1])
await remove_entities(hass, ufp, [camera1])
assert_entity_counts(hass, Platform.CAMERA, 0, 0)
camera1.channels = []
await adopt_devices(hass, ufp, [camera1])
@ -296,7 +296,7 @@ async def test_adopt(hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCa
await hass.async_block_till_done()
assert_entity_counts(hass, Platform.CAMERA, 2, 1)
await remove_entities(hass, [camera1])
await remove_entities(hass, ufp, [camera1])
assert_entity_counts(hass, Platform.CAMERA, 0, 0)
await adopt_devices(hass, ufp, [camera1])
assert_entity_counts(hass, Platform.CAMERA, 2, 1)

2
tests/components/unifiprotect/test_light.py
View file

@ -33,7 +33,7 @@ async def test_light_remove(hass: HomeAssistant, ufp: MockUFPFixture, light: Lig
await init_entry(hass, ufp, [light])
assert_entity_counts(hass, Platform.LIGHT, 1, 1)
await remove_entities(hass, [light])
await remove_entities(hass, ufp, [light])
assert_entity_counts(hass, Platform.LIGHT, 0, 0)
await adopt_devices(hass, ufp, [light])
assert_entity_counts(hass, Platform.LIGHT, 1, 1)

2
tests/components/unifiprotect/test_lock.py
View file

@ -37,7 +37,7 @@ async def test_lock_remove(
await init_entry(hass, ufp, [doorlock])
assert_entity_counts(hass, Platform.LOCK, 1, 1)
await remove_entities(hass, [doorlock])
await remove_entities(hass, ufp, [doorlock])
assert_entity_counts(hass, Platform.LOCK, 0, 0)
await adopt_devices(hass, ufp, [doorlock])
assert_entity_counts(hass, Platform.LOCK, 1, 1)

2
tests/components/unifiprotect/test_media_player.py
View file

@ -41,7 +41,7 @@ async def test_media_player_camera_remove(
await init_entry(hass, ufp, [doorbell])
assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1)
await remove_entities(hass, [doorbell])
await remove_entities(hass, ufp, [doorbell])
assert_entity_counts(hass, Platform.MEDIA_PLAYER, 0, 0)
await adopt_devices(hass, ufp, [doorbell])
assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1)

6
tests/components/unifiprotect/test_number.py
View file

@ -36,7 +36,7 @@ async def test_number_sensor_camera_remove(
await init_entry(hass, ufp, [camera, unadopted_camera])
assert_entity_counts(hass, Platform.NUMBER, 3, 3)
await remove_entities(hass, [camera, unadopted_camera])
await remove_entities(hass, ufp, [camera, unadopted_camera])
assert_entity_counts(hass, Platform.NUMBER, 0, 0)
await adopt_devices(hass, ufp, [camera, unadopted_camera])
assert_entity_counts(hass, Platform.NUMBER, 3, 3)
@ -49,7 +49,7 @@ async def test_number_sensor_light_remove(
await init_entry(hass, ufp, [light])
assert_entity_counts(hass, Platform.NUMBER, 2, 2)
await remove_entities(hass, [light])
await remove_entities(hass, ufp, [light])
assert_entity_counts(hass, Platform.NUMBER, 0, 0)
await adopt_devices(hass, ufp, [light])
assert_entity_counts(hass, Platform.NUMBER, 2, 2)
@ -62,7 +62,7 @@ async def test_number_lock_remove(
await init_entry(hass, ufp, [doorlock])
assert_entity_counts(hass, Platform.NUMBER, 1, 1)
await remove_entities(hass, [doorlock])
await remove_entities(hass, ufp, [doorlock])
assert_entity_counts(hass, Platform.NUMBER, 0, 0)
await adopt_devices(hass, ufp, [doorlock])
assert_entity_counts(hass, Platform.NUMBER, 1, 1)

6
tests/components/unifiprotect/test_select.py
View file

@ -57,7 +57,7 @@ async def test_select_camera_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SELECT, 4, 4)
await remove_entities(hass, [doorbell, unadopted_camera])
await remove_entities(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SELECT, 0, 0)
await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SELECT, 4, 4)
@ -71,7 +71,7 @@ async def test_select_light_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [light])
assert_entity_counts(hass, Platform.SELECT, 2, 2)
await remove_entities(hass, [light])
await remove_entities(hass, ufp, [light])
assert_entity_counts(hass, Platform.SELECT, 0, 0)
await adopt_devices(hass, ufp, [light])
assert_entity_counts(hass, Platform.SELECT, 2, 2)
@ -85,7 +85,7 @@ async def test_select_viewer_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [viewer])
assert_entity_counts(hass, Platform.SELECT, 1, 1)
await remove_entities(hass, [viewer])
await remove_entities(hass, ufp, [viewer])
assert_entity_counts(hass, Platform.SELECT, 0, 0)
await adopt_devices(hass, ufp, [viewer])
assert_entity_counts(hass, Platform.SELECT, 1, 1)

4
tests/components/unifiprotect/test_sensor.py
View file

@ -63,7 +63,7 @@ async def test_sensor_camera_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SENSOR, 25, 13)
await remove_entities(hass, [doorbell, unadopted_camera])
await remove_entities(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SENSOR, 12, 9)
await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SENSOR, 25, 13)
@ -77,7 +77,7 @@ async def test_sensor_sensor_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [sensor_all])
assert_entity_counts(hass, Platform.SENSOR, 22, 14)
await remove_entities(hass, [sensor_all])
await remove_entities(hass, ufp, [sensor_all])
assert_entity_counts(hass, Platform.SENSOR, 12, 9)
await adopt_devices(hass, ufp, [sensor_all])
assert_entity_counts(hass, Platform.SENSOR, 22, 14)

4
tests/components/unifiprotect/test_switch.py
View file

@ -50,7 +50,7 @@ async def test_switch_camera_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SWITCH, 13, 12)
await remove_entities(hass, [doorbell, unadopted_camera])
await remove_entities(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SWITCH, 0, 0)
await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
assert_entity_counts(hass, Platform.SWITCH, 13, 12)
@ -64,7 +64,7 @@ async def test_switch_light_remove(
ufp.api.bootstrap.nvr.system_info.ustorage = None
await init_entry(hass, ufp, [light])
assert_entity_counts(hass, Platform.SWITCH, 2, 1)
await remove_entities(hass, [light])
await remove_entities(hass, ufp, [light])
assert_entity_counts(hass, Platform.SWITCH, 0, 0)
await adopt_devices(hass, ufp, [light])
assert_entity_counts(hass, Platform.SWITCH, 2, 1)

27
tests/components/unifiprotect/utils.py
View file

@ -23,7 +23,7 @@ from pyunifiprotect.test_util.anonymize import random_hex
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant, split_entity_id
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity import EntityDescription
import homeassistant.util.dt as dt_util
@ -89,7 +89,6 @@ def assert_entity_counts(
e for e in entity_registry.entities if split_entity_id(e)[0] == platform.value
]
print(len(entities), total)
assert len(entities) == total
assert len(hass.states.async_all(platform.value)) == enabled
@ -176,28 +175,25 @@ async def init_entry(
async def remove_entities(
hass: HomeAssistant,
ufp: MockUFPFixture,
ufp_devices: list[ProtectAdoptableDeviceModel],
) -> None:
"""Remove all entities for given Protect devices."""
entity_registry = er.async_get(hass)
device_registry = dr.async_get(hass)
for ufp_device in ufp_devices:
if not ufp_device.is_adopted_by_us:
continue
name = ufp_device.display_name.replace(" ", "_").lower()
entity = entity_registry.async_get(f"{Platform.SENSOR}.{name}_uptime")
assert entity is not None
devices = getattr(ufp.api.bootstrap, f"{ufp_device.model.value}s")
del devices[ufp_device.id]
device_id = entity.device_id
for reg in list(entity_registry.entities.values()):
if reg.device_id == device_id:
entity_registry.async_remove(reg.entity_id)
device_registry.async_remove_device(device_id)
mock_msg = Mock()
mock_msg.changed_data = {}
mock_msg.old_obj = ufp_device
mock_msg.new_obj = None
ufp.ws_msg(mock_msg)
await hass.async_block_till_done()
await time_changed(hass, 30)
async def adopt_devices(
@ -214,6 +210,9 @@ async def adopt_devices(
ufp_device.is_adopted_by_other = False
ufp_device.can_adopt = False
devices = getattr(ufp.api.bootstrap, f"{ufp_device.model.value}s")
devices[ufp_device.id] = ufp_device
mock_msg = Mock()
mock_msg.changed_data = {}
mock_msg.new_obj = Event(