Move DPI Group logic to UniFi switch platform (#58761)

* Library has normalized management of DPI apps and groups, move logic to UniFi integration

* Bump dependency to v29

* Use a generator instead of a list - Pylint

* Minor improvements

* Improve doc strings

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(

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"

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."""

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

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

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