Add support for port control in UniFi switch integration (#150152)

This commit is contained in:
Tomeroeni
2025-08-26 17:20:45 +02:00
committed by GitHub
parent a90ac612f0
commit 87f0703be1
3 changed files with 340 additions and 1 deletions

View File

@@ -27,7 +27,10 @@ from aiounifi.interfaces.traffic_rules import TrafficRules
from aiounifi.interfaces.wlans import Wlans
from aiounifi.models.api import ApiItem
from aiounifi.models.client import Client, ClientBlockRequest
from aiounifi.models.device import DeviceSetOutletRelayRequest
from aiounifi.models.device import (
DeviceSetOutletRelayRequest,
DeviceSetPortEnabledRequest,
)
from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest
from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup
from aiounifi.models.event import Event, EventKey
@@ -156,6 +159,14 @@ def async_outlet_switching_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
return outlet.has_relay or outlet.caps in (1, 3)
@callback
def async_port_control_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Determine if a port supports switching."""
port = hub.api.ports[obj_id]
# Only allow switching for physical ports that exist
return port.port_idx is not None
async def async_outlet_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
"""Control outlet relay."""
mac, _, index = obj_id.partition("_")
@@ -174,6 +185,15 @@ async def async_poe_port_control_fn(hub: UnifiHub, obj_id: str, target: bool) ->
hub.queue_poe_port_command(mac, int(index), state)
async def async_port_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
"""Control port enabled state."""
mac, _, index = obj_id.partition("_")
device = hub.api.devices[mac]
await hub.api.request(
DeviceSetPortEnabledRequest.create(device, int(index), target)
)
async def async_port_forward_control_fn(
hub: UnifiHub, obj_id: str, target: bool
) -> None:
@@ -338,6 +358,22 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
supported_fn=lambda hub, obj_id: bool(hub.api.ports[obj_id].port_poe),
unique_id_fn=lambda hub, obj_id: f"poe-{obj_id}",
),
UnifiSwitchEntityDescription[Ports, Port](
key="Port control",
translation_key="port_control",
device_class=SwitchDeviceClass.SWITCH,
entity_category=EntityCategory.CONFIG,
entity_registry_enabled_default=False,
api_handler_fn=lambda api: api.ports,
available_fn=async_device_available_fn,
control_fn=async_port_control_fn,
device_info_fn=async_device_device_info_fn,
is_on_fn=lambda hub, port: bool(port.enabled),
name_fn=lambda port: port.name,
object_fn=lambda api, obj_id: api.ports[obj_id],
supported_fn=async_port_control_supported_fn,
unique_id_fn=lambda hub, obj_id: f"port-{obj_id}",
),
UnifiSwitchEntityDescription[Wlans, Wlan](
key="WLAN control",
translation_key="wlan_control",

View File

@@ -194,6 +194,55 @@
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_1-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
}),
'area_id': None,
'capabilities': None,
'config_entry_id': <ANY>,
'config_subentry_id': <ANY>,
'device_class': None,
'device_id': <ANY>,
'disabled_by': None,
'domain': 'switch',
'entity_category': <EntityCategory.CONFIG: 'config'>,
'entity_id': 'switch.mock_name_port_1',
'has_entity_name': True,
'hidden_by': None,
'icon': None,
'id': <ANY>,
'labels': set({
}),
'name': None,
'options': dict({
}),
'original_device_class': <SwitchDeviceClass.SWITCH: 'switch'>,
'original_icon': None,
'original_name': 'Port 1',
'platform': 'unifi',
'previous_unique_id': None,
'suggested_object_id': None,
'supported_features': 0,
'translation_key': 'port_control',
'unique_id': 'port-10:00:00:00:01:01_1',
'unit_of_measurement': None,
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_1-state]
StateSnapshot({
'attributes': ReadOnlyDict({
'device_class': 'switch',
'friendly_name': 'mock-name Port 1',
}),
'context': <ANY>,
'entity_id': 'switch.mock_name_port_1',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_1_poe-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
@@ -243,6 +292,55 @@
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_2-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
}),
'area_id': None,
'capabilities': None,
'config_entry_id': <ANY>,
'config_subentry_id': <ANY>,
'device_class': None,
'device_id': <ANY>,
'disabled_by': None,
'domain': 'switch',
'entity_category': <EntityCategory.CONFIG: 'config'>,
'entity_id': 'switch.mock_name_port_2',
'has_entity_name': True,
'hidden_by': None,
'icon': None,
'id': <ANY>,
'labels': set({
}),
'name': None,
'options': dict({
}),
'original_device_class': <SwitchDeviceClass.SWITCH: 'switch'>,
'original_icon': None,
'original_name': 'Port 2',
'platform': 'unifi',
'previous_unique_id': None,
'suggested_object_id': None,
'supported_features': 0,
'translation_key': 'port_control',
'unique_id': 'port-10:00:00:00:01:01_2',
'unit_of_measurement': None,
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_2-state]
StateSnapshot({
'attributes': ReadOnlyDict({
'device_class': 'switch',
'friendly_name': 'mock-name Port 2',
}),
'context': <ANY>,
'entity_id': 'switch.mock_name_port_2',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_2_poe-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
@@ -292,6 +390,104 @@
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_3-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
}),
'area_id': None,
'capabilities': None,
'config_entry_id': <ANY>,
'config_subentry_id': <ANY>,
'device_class': None,
'device_id': <ANY>,
'disabled_by': None,
'domain': 'switch',
'entity_category': <EntityCategory.CONFIG: 'config'>,
'entity_id': 'switch.mock_name_port_3',
'has_entity_name': True,
'hidden_by': None,
'icon': None,
'id': <ANY>,
'labels': set({
}),
'name': None,
'options': dict({
}),
'original_device_class': <SwitchDeviceClass.SWITCH: 'switch'>,
'original_icon': None,
'original_name': 'Port 3',
'platform': 'unifi',
'previous_unique_id': None,
'suggested_object_id': None,
'supported_features': 0,
'translation_key': 'port_control',
'unique_id': 'port-10:00:00:00:01:01_3',
'unit_of_measurement': None,
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_3-state]
StateSnapshot({
'attributes': ReadOnlyDict({
'device_class': 'switch',
'friendly_name': 'mock-name Port 3',
}),
'context': <ANY>,
'entity_id': 'switch.mock_name_port_3',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_4-entry]
EntityRegistryEntrySnapshot({
'aliases': set({
}),
'area_id': None,
'capabilities': None,
'config_entry_id': <ANY>,
'config_subentry_id': <ANY>,
'device_class': None,
'device_id': <ANY>,
'disabled_by': None,
'domain': 'switch',
'entity_category': <EntityCategory.CONFIG: 'config'>,
'entity_id': 'switch.mock_name_port_4',
'has_entity_name': True,
'hidden_by': None,
'icon': None,
'id': <ANY>,
'labels': set({
}),
'name': None,
'options': dict({
}),
'original_device_class': <SwitchDeviceClass.SWITCH: 'switch'>,
'original_icon': None,
'original_name': 'Port 4',
'platform': 'unifi',
'previous_unique_id': None,
'suggested_object_id': None,
'supported_features': 0,
'translation_key': 'port_control',
'unique_id': 'port-10:00:00:00:01:01_4',
'unit_of_measurement': None,
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_4-state]
StateSnapshot({
'attributes': ReadOnlyDict({
'device_class': 'switch',
'friendly_name': 'mock-name Port 4',
}),
'context': <ANY>,
'entity_id': 'switch.mock_name_port_4',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'on',
})
# ---
# name: test_entity_and_device_data[site_payload0-wlan_payload0-traffic_rule_payload0-port_forward_payload0-dpi_group_payload0-dpi_app_payload0-device_payload0-client_payload0-config_entry_options0][switch.mock_name_port_4_poe-entry]
EntityRegistryEntrySnapshot({
'aliases': set({

View File

@@ -150,6 +150,7 @@ DEVICE_1 = {
"portconf_id": "1a1",
"port_poe": True,
"up": True,
"enable": True,
},
{
"media": "GE",
@@ -164,6 +165,7 @@ DEVICE_1 = {
"portconf_id": "1a2",
"port_poe": True,
"up": True,
"enable": True,
},
{
"media": "GE",
@@ -178,6 +180,7 @@ DEVICE_1 = {
"portconf_id": "1a3",
"port_poe": False,
"up": True,
"enable": True,
},
{
"media": "GE",
@@ -192,6 +195,7 @@ DEVICE_1 = {
"portconf_id": "1a4",
"port_poe": True,
"up": True,
"enable": True,
},
],
"state": 1,
@@ -1727,6 +1731,7 @@ async def test_port_forwarding_switches(
"portconf_id": "1a1",
"port_poe": True,
"up": True,
"enable": True,
},
],
},
@@ -1783,6 +1788,7 @@ async def test_hub_state_change(
entity_ids = (
"switch.block_client_2",
"switch.mock_name_port_1_poe",
"switch.mock_name_port_1",
"switch.plug_outlet_1",
"switch.block_media_streaming",
"switch.unifi_network_plex",
@@ -1802,3 +1808,104 @@ async def test_hub_state_change(
await mock_websocket_state.reconnect()
for entity_id in entity_ids:
assert hass.states.get(entity_id).state == STATE_ON
@pytest.mark.parametrize("device_payload", [[DEVICE_1]])
async def test_port_control_switches(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker,
config_entry_setup: MockConfigEntry,
mock_websocket_message: WebsocketMessageMock,
device_payload: list[dict[str, Any]],
) -> None:
"""Test port control entities work."""
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
ent_reg_entry = entity_registry.async_get("switch.mock_name_port_1")
assert (
ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
) # ✅ Disabled by default
# Enable entity
entity_registry.async_update_entity(
entity_id="switch.mock_name_port_1", disabled_by=None
)
entity_registry.async_update_entity(
entity_id="switch.mock_name_port_2", disabled_by=None
)
async_fire_time_changed(
hass,
dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1),
)
await hass.async_block_till_done()
# Validate state object
assert hass.states.get("switch.mock_name_port_1").state == STATE_ON
# Update state object - disable port via port_overrides
device_1 = deepcopy(device_payload[0])
device_1["port_table"][0]["enable"] = False
mock_websocket_message(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.mock_name_port_1").state == STATE_OFF
# Turn off port
aioclient_mock.clear_requests()
aioclient_mock.put(
f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/rest/device/mock-id",
)
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_off",
{"entity_id": "switch.mock_name_port_1"},
blocking=True,
)
async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=5))
await hass.async_block_till_done()
assert aioclient_mock.call_count == 1
assert aioclient_mock.mock_calls[0][2] == {
"port_overrides": [{"enable": False, "port_idx": 1, "portconf_id": "1a1"}]
}
# Turn on port
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_on",
{"entity_id": "switch.mock_name_port_1"},
blocking=True,
)
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_off",
{"entity_id": "switch.mock_name_port_2"},
blocking=True,
)
async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=5))
await hass.async_block_till_done()
assert aioclient_mock.call_count == 3
assert aioclient_mock.mock_calls[1][2] == {
"port_overrides": [
{"port_idx": 1, "enable": True, "portconf_id": "1a1"},
]
}
assert aioclient_mock.mock_calls[2][2] == {
"port_overrides": [
{"port_idx": 2, "enable": False, "portconf_id": "1a2"},
]
}
# Device gets disabled
device_1["disabled"] = True
mock_websocket_message(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.mock_name_port_1").state == STATE_UNAVAILABLE
# Device gets re-enabled
device_1["disabled"] = False
mock_websocket_message(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.mock_name_port_1").state == STATE_OFF