Expose UniFi Policy Engine rules as switches (#169675)

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

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,

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

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