Expose UniFi Policy Engine rules as switches

homeassistant/components/unifi/hub/entity_loader.py

@@ -49,6 +49,9 @@ def __init__(self, hub: UnifiHub) -> None:
         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/manifest.json

@@ -9,5 +9,5 @@
   "iot_class": "local_push",
   "loggers": ["aiounifi"],
   "quality_scale": "silver",
-  "requirements": ["aiounifi==90"]
+  "requirements": ["aiounifi==91"]
 }

homeassistant/components/unifi/switch.py

@@ -5,6 +5,7 @@
 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.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_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,
+    ObjectOrientedNetworkConfigUpdateRequest,
+)
 from aiounifi.models.outlet import Outlet
 from aiounifi.models.port import Port
 from aiounifi.models.port_forward import PortForward, PortForwardEnableRequest
@@ -153,6 +161,16 @@ 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].raw
+    await hub.api.request(
+        ObjectOrientedNetworkConfigUpdateRequest.create(config, target)
+    )
+
+
 @callback
 def async_outlet_switching_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
     """Determine if an outlet supports switching."""
@@ -284,6 +302,20 @@ class UnifiSwitchEntityDescription[HandlerT: APIHandler, ApiItemT: ApiItem](
         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],
+        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 mock_get_request(path: str, payload: list[dict[str, Any]]) -> None:
         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,23 @@
     },
 }
 
+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"},
+        },
+    },
+}
+
 
 @pytest.mark.parametrize(
     "config_entry_options", [{CONF_BLOCK_CLIENT: [BLOCKED["mac"]]}]
@@ -1351,6 +1368,65 @@ 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])]
+)
+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."""
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
+
+    # Validate state object
+    assert (
+        hass.states.get("switch.unifi_network_nintendo_switch_block_internet").state
+        == STATE_ON
+    )
+
+    config = deepcopy(object_oriented_network_config_payload[0])
+
+    # Disable Policy Engine rule
+    aioclient_mock.put(
+        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']}",
+    )
+
+    call_count = aioclient_mock.call_count
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        "turn_off",
+        {"entity_id": "switch.unifi_network_nintendo_switch_block_internet"},
+        blocking=True,
+    )
+    # Updating the value for Policy Engine rules will make another call to retrieve the values
+    assert aioclient_mock.call_count == call_count + 2
+    expected_disable_call = deepcopy(config)
+    expected_disable_call["enabled"] = False
+
+    assert aioclient_mock.mock_calls[call_count][2] == expected_disable_call
+
+    call_count = aioclient_mock.call_count
+
+    # Enable Policy Engine rule
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        "turn_on",
+        {"entity_id": "switch.unifi_network_nintendo_switch_block_internet"},
+        blocking=True,
+    )
+
+    expected_enable_call = deepcopy(config)
+    expected_enable_call["enabled"] = True
+
+    assert aioclient_mock.call_count == call_count + 1
+    assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call
+
+
 @pytest.mark.parametrize(
     ("device_payload", "entity_id", "outlet_index", "expected_switches"),
     [