From c6d34dd31539aa2eff48033ff74c026e5da89db0 Mon Sep 17 00:00:00 2001 From: Ryan Ray Date: Tue, 18 Aug 2026 01:25:30 -0500 Subject: [PATCH] Expose UniFi Policy Engine rules as switches (#169675) --- .../components/unifi/hub/entity_loader.py | 3 + homeassistant/components/unifi/switch.py | 46 +++++++++ tests/components/unifi/conftest.py | 11 +++ tests/components/unifi/test_switch.py | 96 +++++++++++++++++++ 4 files changed, 156 insertions(+) diff --git a/homeassistant/components/unifi/hub/entity_loader.py b/homeassistant/components/unifi/hub/entity_loader.py index abc0ece041d8..9e70d51cc70b 100644 --- a/homeassistant/components/unifi/hub/entity_loader.py +++ b/homeassistant/components/unifi/hub/entity_loader.py @@ -51,6 +51,9 @@ class UnifiEntityLoader: self.wireless_clients = hub.hass.data[UNIFI_WIRELESS_CLIENTS] self._polling_coordinators: dict[int, UnifiDataUpdateCoordinator] = { + id(hub.api.object_oriented_network_configs): UnifiDataUpdateCoordinator( + hub, hub.api.object_oriented_network_configs + ), id(hub.api.traffic_rules): UnifiDataUpdateCoordinator( hub, hub.api.traffic_rules ), diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 4677c023c347..0173a396fd15 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -5,6 +5,7 @@ Support for controlling network access of clients selected in option flow. Support for controlling deep packet inspection (DPI) restriction groups. Support for controlling WLAN availability. Support for controlling zone based traffic rules. +Support for controlling Policy Engine rules. """ import asyncio @@ -17,6 +18,9 @@ from aiounifi.interfaces.api_handlers import APIHandler, ItemEvent from aiounifi.interfaces.clients import Clients from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups from aiounifi.interfaces.firewall_policies import FirewallPolicies +from aiounifi.interfaces.object_oriented_network_configs import ( + ObjectOrientedNetworkConfigs, +) from aiounifi.interfaces.outlets import Outlets from aiounifi.interfaces.port_forwarding import PortForwarding from aiounifi.interfaces.ports import Ports @@ -33,6 +37,10 @@ from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup from aiounifi.models.event import Event, EventKey from aiounifi.models.firewall_policy import FirewallPolicy, FirewallPolicyUpdateRequest +from aiounifi.models.object_oriented_network_config import ( + ObjectOrientedNetworkConfig, + ObjectOrientedNetworkInternetMode, +) from aiounifi.models.outlet import Outlet from aiounifi.models.port import Port from aiounifi.models.port_forward import PortForward, PortForwardEnableRequest @@ -152,6 +160,29 @@ def async_firewall_policy_supported_fn(hub: UnifiHub, obj_id: str) -> bool: return not policy.predefined +async def async_object_oriented_network_config_control_fn( + hub: UnifiHub, obj_id: str, target: bool +) -> None: + """Control Policy Engine rule state.""" + config = hub.api.object_oriented_network_configs[obj_id] + await hub.api.object_oriented_network_configs.save(config, target) + + +@callback +def async_object_oriented_network_config_supported_fn( + hub: UnifiHub, obj_id: str +) -> bool: + """Check if Policy Engine rule can be controlled as a switch.""" + config = hub.api.object_oriented_network_configs[obj_id] + secure = config.secure + return ( + secure.available + and secure.enabled + and secure.internet is not None + and secure.internet.mode is ObjectOrientedNetworkInternetMode.TURN_OFF_INTERNET + ) + + @callback def async_outlet_switching_supported_fn(hub: UnifiHub, obj_id: str) -> bool: """Determine if an outlet supports switching.""" @@ -283,6 +314,21 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( unique_id_fn=lambda hub, obj_id: f"firewall_policy-{obj_id}", supported_fn=async_firewall_policy_supported_fn, ), + UnifiSwitchEntityDescription[ + ObjectOrientedNetworkConfigs, ObjectOrientedNetworkConfig + ]( + key="Policy Engine rule control", + device_class=SwitchDeviceClass.SWITCH, + entity_category=EntityCategory.CONFIG, + api_handler_fn=lambda api: api.object_oriented_network_configs, + control_fn=async_object_oriented_network_config_control_fn, + device_info_fn=async_unifi_network_device_info_fn, + is_on_fn=lambda hub, config: config.enabled, + name_fn=lambda config: config.name, + object_fn=lambda api, obj_id: api.object_oriented_network_configs[obj_id], + supported_fn=async_object_oriented_network_config_supported_fn, + unique_id_fn=lambda hub, obj_id: f"object_oriented_network_config-{obj_id}", + ), UnifiSwitchEntityDescription[Outlets, Outlet]( key="Outlet control", device_class=SwitchDeviceClass.OUTLET, diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index dc08dedf1171..a33d3b536eff 100644 --- a/tests/components/unifi/conftest.py +++ b/tests/components/unifi/conftest.py @@ -179,6 +179,7 @@ def fixture_request( dpi_app_payload: list[dict[str, Any]], dpi_group_payload: list[dict[str, Any]], firewall_policy_payload: list[dict[str, Any]], + object_oriented_network_config_payload: list[dict[str, Any]], port_forward_payload: list[dict[str, Any]], traffic_rule_payload: list[dict[str, Any]], traffic_route_payload: list[dict[str, Any]], @@ -221,6 +222,10 @@ def fixture_request( mock_get_request( f"/v2/api/site/{site_id}/firewall-policies", firewall_policy_payload ) + mock_get_request( + f"/v2/api/site/{site_id}/object-oriented-network-configs", + object_oriented_network_config_payload, + ) mock_get_request(f"/api/s/{site_id}/rest/portforward", port_forward_payload) mock_get_request(f"/api/s/{site_id}/stat/sysinfo", system_information_payload) mock_get_request(f"/api/s/{site_id}/rest/wlanconf", wlan_payload) @@ -269,6 +274,12 @@ def firewall_policy_payload_data() -> list[dict[str, Any]]: return [] +@pytest.fixture(name="object_oriented_network_config_payload") +def object_oriented_network_config_payload_data() -> list[dict[str, Any]]: + """Object-oriented network config data.""" + return [] + + @pytest.fixture(name="port_forward_payload") def fixture_port_forward_data() -> list[dict[str, Any]]: """Port forward data.""" diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 05b9db5ddb58..e1c980a0e210 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -872,6 +872,33 @@ FIREWALL_POLICY = { }, } +OBJECT_ORIENTED_NETWORK_CONFIG = { + "id": "69f6b0a5e0e3ee2d4614cb5c", + "enabled": True, + "name": "Nintendo Switch - Block Internet", + "target_type": "CLIENTS", + "targets": [CLIENT_1["mac"]], + "qos": {"enabled": False}, + "route": {"enabled": False}, + "secure": { + "enabled": True, + "internet": { + "mode": "TURN_OFF_INTERNET", + "schedule": {"mode": "ALWAYS"}, + }, + }, +} + +OBJECT_ORIENTED_NETWORK_ROUTE_CONFIG = { + "id": "69f6b0eae0e3ee2d4614cb91", + "enabled": True, + "name": "VPN traffic route", + "target_type": "NETWORKS", + "targets": ["6060b00f45de3905133cea14"], + "route": {"enabled": True}, + "secure": None, +} + @pytest.mark.parametrize( "config_entry_options", [{CONF_BLOCK_CLIENT: [BLOCKED["mac"]]}] @@ -1353,6 +1380,75 @@ async def test_firewall_policies( assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call +@pytest.mark.parametrize( + ("object_oriented_network_config_payload"), + [([OBJECT_ORIENTED_NETWORK_CONFIG, OBJECT_ORIENTED_NETWORK_ROUTE_CONFIG])], +) +async def test_object_oriented_network_configs( + hass: HomeAssistant, + aioclient_mock: AiohttpClientMocker, + config_entry_setup: MockConfigEntry, + object_oriented_network_config_payload: list[dict[str, Any]], +) -> None: + """Test control of UniFi Policy Engine rules.""" + entity_id = "switch.unifi_network_nintendo_switch_block_internet" + assert hass.states.get("switch.unifi_network_vpn_traffic_route") is None + + state = hass.states.get(entity_id) + assert state is not None + assert state.state == STATE_ON + + config = deepcopy(object_oriented_network_config_payload[0]) + config_url = ( + f"https://{config_entry_setup.data[CONF_HOST]}:1234" + f"/v2/api/site/{config_entry_setup.data[CONF_SITE_ID]}" + f"/object-oriented-network-config/{config['id']}" + ) + + aioclient_mock.put(config_url) + + call_count = aioclient_mock.call_count + + await hass.services.async_call( + SWITCH_DOMAIN, + "turn_off", + {"entity_id": entity_id}, + blocking=True, + ) + expected_disable_call = deepcopy(config) + expected_disable_call["enabled"] = False + + assert ( + "put", + config_url, + expected_disable_call, + ) in ( + (method, str(url), data) + for method, url, data, _headers in aioclient_mock.mock_calls[call_count:] + ) + + call_count = aioclient_mock.call_count + + await hass.services.async_call( + SWITCH_DOMAIN, + "turn_on", + {"entity_id": entity_id}, + blocking=True, + ) + + expected_enable_call = deepcopy(config) + expected_enable_call["enabled"] = True + + assert ( + "put", + config_url, + expected_enable_call, + ) in ( + (method, str(url), data) + for method, url, data, _headers in aioclient_mock.mock_calls[call_count:] + ) + + @pytest.mark.parametrize( ("device_payload", "entity_id", "outlet_index", "expected_switches"), [