From 87f0703be17a20f0f87a1d7d76ca0781173555ab Mon Sep 17 00:00:00 2001 From: Tomeroeni <30298350+Tomeroeni@users.noreply.github.com> Date: Tue, 26 Aug 2025 17:20:45 +0200 Subject: [PATCH] Add support for port control in UniFi switch integration (#150152) --- homeassistant/components/unifi/switch.py | 38 +++- .../unifi/snapshots/test_switch.ambr | 196 ++++++++++++++++++ tests/components/unifi/test_switch.py | 107 ++++++++++ 3 files changed, 340 insertions(+), 1 deletion(-) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 1ca409bec77..b9fbf48cf49 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -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", diff --git a/tests/components/unifi/snapshots/test_switch.ambr b/tests/components/unifi/snapshots/test_switch.ambr index 017fe237025..4fabff5d278 100644 --- a/tests/components/unifi/snapshots/test_switch.ambr +++ b/tests/components/unifi/snapshots/test_switch.ambr @@ -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': , + 'config_subentry_id': , + 'device_class': None, + 'device_id': , + 'disabled_by': None, + 'domain': 'switch', + 'entity_category': , + 'entity_id': 'switch.mock_name_port_1', + 'has_entity_name': True, + 'hidden_by': None, + 'icon': None, + 'id': , + 'labels': set({ + }), + 'name': None, + 'options': dict({ + }), + 'original_device_class': , + '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': , + 'entity_id': 'switch.mock_name_port_1', + 'last_changed': , + 'last_reported': , + 'last_updated': , + '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': , + 'config_subentry_id': , + 'device_class': None, + 'device_id': , + 'disabled_by': None, + 'domain': 'switch', + 'entity_category': , + 'entity_id': 'switch.mock_name_port_2', + 'has_entity_name': True, + 'hidden_by': None, + 'icon': None, + 'id': , + 'labels': set({ + }), + 'name': None, + 'options': dict({ + }), + 'original_device_class': , + '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': , + 'entity_id': 'switch.mock_name_port_2', + 'last_changed': , + 'last_reported': , + 'last_updated': , + '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': , + 'config_subentry_id': , + 'device_class': None, + 'device_id': , + 'disabled_by': None, + 'domain': 'switch', + 'entity_category': , + 'entity_id': 'switch.mock_name_port_3', + 'has_entity_name': True, + 'hidden_by': None, + 'icon': None, + 'id': , + 'labels': set({ + }), + 'name': None, + 'options': dict({ + }), + 'original_device_class': , + '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': , + 'entity_id': 'switch.mock_name_port_3', + 'last_changed': , + 'last_reported': , + 'last_updated': , + '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': , + 'config_subentry_id': , + 'device_class': None, + 'device_id': , + 'disabled_by': None, + 'domain': 'switch', + 'entity_category': , + 'entity_id': 'switch.mock_name_port_4', + 'has_entity_name': True, + 'hidden_by': None, + 'icon': None, + 'id': , + 'labels': set({ + }), + 'name': None, + 'options': dict({ + }), + 'original_device_class': , + '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': , + 'entity_id': 'switch.mock_name_port_4', + 'last_changed': , + 'last_reported': , + 'last_updated': , + '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({ diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index c14ecbc0b06..442bc4f83e6 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -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