1
0
mirror of https://github.com/home-assistant/core.git synced 2025-12-24 12:59:34 +00:00

Add DPI Restriction switch to UniFi integration (#42499)

* initial implementation for controlling DPI restrictions

* address PR review comments and add DataUpdateCoordinator

* fix existing tests against new lib version

* add tests for DPI switches

* bump aiounifi

* listen to events instead of polling

* fix tests

* remove useless test

* bump aiounifi

* rename device to UniFi Controller per PR feedback
This commit is contained in:
Jason Hunter
2020-11-03 02:36:37 -05:00
committed by GitHub
parent aab0ff2ea5
commit 5a4c1dbcc4
14 changed files with 309 additions and 42 deletions

View File

@@ -9,6 +9,8 @@ from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.components.unifi.const import (
CONF_BLOCK_CLIENT,
CONF_DPI_RESTRICTIONS,
CONF_POE_CLIENTS,
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
DOMAIN as UNIFI_DOMAIN,
@@ -251,6 +253,35 @@ EVENT_CLIENT_2_CONNECTED = {
}
DPI_GROUPS = [
{
"_id": "5ba29dd8e3c58f026e9d7c4a",
"attr_no_delete": True,
"attr_hidden_id": "Default",
"name": "Default",
"site_id": "name",
},
{
"_id": "5f976f4ae3c58f018ec7dff6",
"name": "Block Media Streaming",
"site_id": "name",
"dpiapp_ids": ["5f976f62e3c58f018ec7e17d"],
},
]
DPI_APPS = [
{
"_id": "5f976f62e3c58f018ec7e17d",
"apps": [],
"blocked": True,
"cats": ["4"],
"enabled": True,
"log": True,
"site_id": "name",
}
]
async def test_platform_manually_configured(hass):
"""Test that we do not discover anything or try to set up a controller."""
assert (
@@ -266,10 +297,14 @@ async def test_no_clients(hass):
"""Test the update_clients function when no clients are found."""
controller = await setup_unifi_integration(
hass,
options={CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False},
options={
CONF_TRACK_CLIENTS: False,
CONF_TRACK_DEVICES: False,
CONF_DPI_RESTRICTIONS: False,
},
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@@ -282,7 +317,7 @@ async def test_controller_not_client(hass):
devices_response=[DEVICE_1],
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
cloudkey = hass.states.get("switch.cloud_key")
assert cloudkey is None
@@ -300,7 +335,7 @@ async def test_not_admin(hass):
devices_response=[DEVICE_1],
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@@ -316,10 +351,12 @@ async def test_switches(hass):
clients_response=[CLIENT_1, CLIENT_4],
devices_response=[DEVICE_1],
clients_all_response=[BLOCKED, UNBLOCKED, CLIENT_1],
dpigroup_response=DPI_GROUPS,
dpiapp_response=DPI_APPS,
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 4
switch_1 = hass.states.get("switch.poe_client_1")
assert switch_1 is not None
@@ -340,11 +377,15 @@ async def test_switches(hass):
assert unblocked is not None
assert unblocked.state == "on"
dpi_switch = hass.states.get("switch.block_media_streaming")
assert dpi_switch is not None
assert dpi_switch.state == "on"
await hass.services.async_call(
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 5
assert controller.mock_requests[4] == {
assert len(controller.mock_requests) == 7
assert controller.mock_requests[6] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
"method": "post",
"path": "/cmd/stamgr",
@@ -353,13 +394,39 @@ async def test_switches(hass):
await hass.services.async_call(
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 6
assert controller.mock_requests[5] == {
assert len(controller.mock_requests) == 8
assert controller.mock_requests[7] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
"method": "post",
"path": "/cmd/stamgr",
}
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_off",
{"entity_id": "switch.block_media_streaming"},
blocking=True,
)
assert len(controller.mock_requests) == 9
assert controller.mock_requests[8] == {
"json": {"enabled": False},
"method": "put",
"path": "/rest/dpiapp/5f976f62e3c58f018ec7e17d",
}
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_on",
{"entity_id": "switch.block_media_streaming"},
blocking=True,
)
assert len(controller.mock_requests) == 10
assert controller.mock_requests[9] == {
"json": {"enabled": True},
"method": "put",
"path": "/rest/dpiapp/5f976f62e3c58f018ec7e17d",
}
async def test_remove_switches(hass):
"""Test the update_items function with some clients."""
@@ -443,8 +510,8 @@ async def test_block_switches(hass):
await hass.services.async_call(
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 5
assert controller.mock_requests[4] == {
assert len(controller.mock_requests) == 7
assert controller.mock_requests[6] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
"method": "post",
"path": "/cmd/stamgr",
@@ -453,8 +520,8 @@ async def test_block_switches(hass):
await hass.services.async_call(
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 6
assert controller.mock_requests[5] == {
assert len(controller.mock_requests) == 8
assert controller.mock_requests[7] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
"method": "post",
"path": "/cmd/stamgr",
@@ -469,10 +536,11 @@ async def test_new_client_discovered_on_block_control(hass):
CONF_BLOCK_CLIENT: [BLOCKED["mac"]],
CONF_TRACK_CLIENTS: False,
CONF_TRACK_DEVICES: False,
CONF_DPI_RESTRICTIONS: False,
},
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
blocked = hass.states.get("switch.block_client_1")
@@ -541,6 +609,30 @@ async def test_option_block_clients(hass):
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
async def test_option_remove_switches(hass):
"""Test removal of DPI switch when options updated."""
controller = await setup_unifi_integration(
hass,
options={
CONF_TRACK_CLIENTS: False,
CONF_TRACK_DEVICES: False,
},
clients_response=[CLIENT_1],
devices_response=[DEVICE_1],
dpigroup_response=DPI_GROUPS,
dpiapp_response=DPI_APPS,
)
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
# Disable DPI Switches
hass.config_entries.async_update_entry(
controller.config_entry,
options={CONF_DPI_RESTRICTIONS: False, CONF_POE_CLIENTS: False},
)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
async def test_new_client_discovered_on_poe_control(hass):
"""Test if 2nd update has a new client."""
controller = await setup_unifi_integration(
@@ -550,7 +642,7 @@ async def test_new_client_discovered_on_poe_control(hass):
devices_response=[DEVICE_1],
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
controller.api.websocket._data = {
@@ -576,9 +668,9 @@ async def test_new_client_discovered_on_poe_control(hass):
await hass.services.async_call(
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.poe_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 5
assert len(controller.mock_requests) == 7
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
assert controller.mock_requests[4] == {
assert controller.mock_requests[6] == {
"json": {
"port_overrides": [{"port_idx": 1, "portconf_id": "1a1", "poe_mode": "off"}]
},
@@ -589,8 +681,8 @@ async def test_new_client_discovered_on_poe_control(hass):
await hass.services.async_call(
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.poe_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 6
assert controller.mock_requests[4] == {
assert len(controller.mock_requests) == 8
assert controller.mock_requests[7] == {
"json": {
"port_overrides": [
{"port_idx": 1, "portconf_id": "1a1", "poe_mode": "auto"}
@@ -613,7 +705,7 @@ async def test_ignore_multiple_poe_clients_on_same_port(hass):
devices_response=[DEVICE_1],
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3
switch_1 = hass.states.get("switch.poe_client_1")
@@ -664,7 +756,7 @@ async def test_restoring_client(hass):
clients_all_response=[CLIENT_1],
)
assert len(controller.mock_requests) == 4
assert len(controller.mock_requests) == 6
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
device_1 = hass.states.get("switch.client_1")