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
This commit is contained in:
Robert Svensson
GitHub
parent
b71a22557d
commit
e37456fb36
6 changed files with 190 additions and 17 deletions
|
|
@ -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})
|
||||
|
||||
elif DATA_DPI_GROUP_REMOVED in data:
|
||||
async_dispatcher_send(
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
"config_flow": true,
|
||||
"documentation": "https://www.home-assistant.io/integrations/unifi",
|
||||
"requirements": [
|
||||
"aiounifi==28"
|
||||
"aiounifi==29"
|
||||
],
|
||||
"codeowners": [
|
||||
"@Kane610"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue