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

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

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': <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({

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