diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 209ff9980c9..add1b7cdd45 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -246,11 +246,7 @@ class UniFiController: ) elif DATA_DPI_GROUP in data: - for key in data[DATA_DPI_GROUP]: - if self.api.dpi_groups[key].dpiapp_ids: - async_dispatcher_send(self.hass, self.signal_update) - else: - async_dispatcher_send(self.hass, self.signal_remove, {key}) + async_dispatcher_send(self.hass, self.signal_update) elif DATA_DPI_GROUP_REMOVED in data: async_dispatcher_send( diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index 7dbd86c928a..4dda43da9d0 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -4,7 +4,7 @@ "config_flow": true, "documentation": "https://www.home-assistant.io/integrations/unifi", "requirements": [ - "aiounifi==28" + "aiounifi==29" ], "codeowners": [ "@Kane610" diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index fdb4c6af3da..9ba9ed6dd0f 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -4,6 +4,8 @@ Support for controlling power supply of clients which are powered over Ethernet Support for controlling network access of clients selected in option flow. Support for controlling deep packet inspection (DPI) restriction groups. """ + +import asyncio from typing import Any from aiounifi.api import SOURCE_EVENT @@ -332,11 +334,57 @@ class UniFiDPIRestrictionSwitch(UniFiBase, SwitchEntity): _attr_entity_category = EntityCategory.CONFIG + def __init__(self, dpi_group, controller): + """Set up dpi switch.""" + super().__init__(dpi_group, controller) + + self._is_enabled = self.calculate_enabled() + self._known_app_ids = dpi_group.dpiapp_ids + @property def key(self) -> Any: """Return item key.""" return self._item.id + async def async_added_to_hass(self) -> None: + """Register callback to known apps.""" + await super().async_added_to_hass() + + apps = self.controller.api.dpi_apps + for app_id in self._item.dpiapp_ids: + apps[app_id].register_callback(self.async_update_callback) + + async def async_will_remove_from_hass(self) -> None: + """Remove registered callbacks.""" + apps = self.controller.api.dpi_apps + for app_id in self._item.dpiapp_ids: + apps[app_id].remove_callback(self.async_update_callback) + + await super().async_will_remove_from_hass() + + @callback + def async_update_callback(self) -> None: + """Update the DPI switch state. + + Remove entity when no apps are paired with group. + Register callbacks to new apps. + Calculate and update entity state if it has changed. + """ + if not self._item.dpiapp_ids: + self.hass.loop.create_task(self.remove_item({self.key})) + return + + if self._known_app_ids != self._item.dpiapp_ids: + self._known_app_ids = self._item.dpiapp_ids + + apps = self.controller.api.dpi_apps + for app_id in self._item.dpiapp_ids: + apps[app_id].register_callback(self.async_update_callback) + + if (enabled := self.calculate_enabled()) != self._is_enabled: + self._is_enabled = enabled + super().async_update_callback() + @property def unique_id(self): """Return a unique identifier for this switch.""" @@ -344,28 +392,46 @@ class UniFiDPIRestrictionSwitch(UniFiBase, SwitchEntity): @property def name(self) -> str: - """Return the name of the client.""" + """Return the name of the DPI group.""" return self._item.name @property def icon(self): """Return the icon to use in the frontend.""" - if self._item.enabled: + if self._is_enabled: return "mdi:network" return "mdi:network-off" + def calculate_enabled(self) -> bool: + """Calculate if all apps are enabled.""" + return all( + self.controller.api.dpi_apps[app_id].enabled + for app_id in self._item.dpiapp_ids + if app_id in self.controller.api.dpi_apps + ) + @property def is_on(self): - """Return true if client is allowed to connect.""" - return self._item.enabled + """Return true if DPI group app restriction is enabled.""" + return self._is_enabled async def async_turn_on(self, **kwargs): - """Turn on connectivity for client.""" - await self.controller.api.dpi_groups.async_enable(self._item) + """Restrict access of apps related to DPI group.""" + return await asyncio.gather( + *[ + self.controller.api.dpi_apps.async_enable(app_id) + for app_id in self._item.dpiapp_ids + ] + ) async def async_turn_off(self, **kwargs): - """Turn off connectivity for client.""" - await self.controller.api.dpi_groups.async_disable(self._item) + """Remove restriction of apps related to DPI group.""" + return await asyncio.gather( + *[ + self.controller.api.dpi_apps.async_disable(app_id) + for app_id in self._item.dpiapp_ids + ] + ) async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" diff --git a/requirements_all.txt b/requirements_all.txt index a9c1e37b038..3870d3acaf3 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -272,7 +272,7 @@ aiosyncthing==0.5.1 aiotractive==0.5.2 # homeassistant.components.unifi -aiounifi==28 +aiounifi==29 # homeassistant.components.vlc_telnet aiovlc==0.1.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 9715f463f55..cc081fd9ad6 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -204,7 +204,7 @@ aiosyncthing==0.5.1 aiotractive==0.5.2 # homeassistant.components.unifi -aiounifi==28 +aiounifi==29 # homeassistant.components.vlc_telnet aiovlc==0.1.0 diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index a5ca1c0ee6a..73226a7f647 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -288,7 +288,42 @@ DPI_GROUP_REMOVED_EVENT = { "data": [ { "_id": "5f976f4ae3c58f018ec7dff6", - "name": "dpi group", + "name": "Block Media Streaming", + "site_id": "name", + "dpiapp_ids": [], + } + ], +} + +DPI_GROUP_CREATED_EVENT = { + "meta": {"rc": "ok", "message": "dpigroup:add"}, + "data": [ + { + "name": "Block Media Streaming", + "site_id": "name", + "_id": "5f976f4ae3c58f018ec7dff6", + } + ], +} + +DPI_GROUP_ADDED_APP = { + "meta": {"rc": "ok", "message": "dpigroup:sync"}, + "data": [ + { + "_id": "5f976f4ae3c58f018ec7dff6", + "name": "Block Media Streaming", + "site_id": "name", + "dpiapp_ids": ["5f976f62e3c58f018ec7e17d"], + } + ], +} + +DPI_GROUP_REMOVE_APP = { + "meta": {"rc": "ok", "message": "dpigroup:sync"}, + "data": [ + { + "_id": "5f976f4ae3c58f018ec7dff6", + "name": "Block Media Streaming", "site_id": "name", "dpiapp_ids": [], } @@ -599,6 +634,82 @@ async def test_dpi_switches(hass, aioclient_mock, mock_unifi_websocket): assert hass.states.get("switch.block_media_streaming").state == STATE_OFF + mock_unifi_websocket(data=DPI_GROUP_REMOVE_APP) + await hass.async_block_till_done() + await hass.async_block_till_done() + await hass.async_block_till_done() + + assert hass.states.get("switch.block_media_streaming") is None + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 + + +async def test_dpi_switches_add_second_app(hass, aioclient_mock, mock_unifi_websocket): + """Test the update_items function with some clients.""" + await setup_unifi_integration( + hass, + aioclient_mock, + dpigroup_response=DPI_GROUPS, + dpiapp_response=DPI_APPS, + ) + + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 + assert hass.states.get("switch.block_media_streaming").state == STATE_ON + + second_app_event = { + "meta": {"rc": "ok", "message": "dpiapp:add"}, + "data": [ + { + "apps": [524292], + "blocked": False, + "cats": [], + "enabled": False, + "log": False, + "site_id": "name", + "_id": "61783e89c1773a18c0c61f00", + } + ], + } + mock_unifi_websocket(data=second_app_event) + await hass.async_block_till_done() + + assert hass.states.get("switch.block_media_streaming").state == STATE_ON + + add_second_app_to_group = { + "meta": {"rc": "ok", "message": "dpigroup:sync"}, + "data": [ + { + "_id": "5f976f4ae3c58f018ec7dff6", + "name": "Block Media Streaming", + "site_id": "name", + "dpiapp_ids": ["5f976f62e3c58f018ec7e17d", "61783e89c1773a18c0c61f00"], + } + ], + } + + mock_unifi_websocket(data=add_second_app_to_group) + await hass.async_block_till_done() + + assert hass.states.get("switch.block_media_streaming").state == STATE_OFF + + second_app_event_enabled = { + "meta": {"rc": "ok", "message": "dpiapp:sync"}, + "data": [ + { + "apps": [524292], + "blocked": False, + "cats": [], + "enabled": True, + "log": False, + "site_id": "name", + "_id": "61783e89c1773a18c0c61f00", + } + ], + } + mock_unifi_websocket(data=second_app_event_enabled) + await hass.async_block_till_done() + + assert hass.states.get("switch.block_media_streaming").state == STATE_ON + async def test_new_client_discovered_on_block_control( hass, aioclient_mock, mock_unifi_websocket