Create switches for controlling policy-based routes (#134473)

Create switches for controlling policy-based routes (aka "traffic routes" in the Unifi API).

homeassistant/components/unifi/hub/entity_loader.py

@@ -47,9 +47,13 @@ class UnifiEntityLoader:
             hub.api.sites.update,
             hub.api.system_information.update,
             hub.api.traffic_rules.update,
+            hub.api.traffic_routes.update,
             hub.api.wlans.update,
         )
-        self.polling_api_updaters = (hub.api.traffic_rules.update,)
+        self.polling_api_updaters = (
+            hub.api.traffic_rules.update,
+            hub.api.traffic_routes.update,
+        )
         self.wireless_clients = hub.hass.data[UNIFI_WIRELESS_CLIENTS]
 
         self._dataUpdateCoordinator = DataUpdateCoordinator(

homeassistant/components/unifi/icons.json

@@ -61,6 +61,9 @@
       "traffic_rule_control": {
         "default": "mdi:security-network"
       },
+      "traffic_route_control": {
+        "default": "mdi:routes"
+      },
       "poe_port_control": {
         "default": "mdi:ethernet",
         "state": {

homeassistant/components/unifi/switch.py

@@ -20,6 +20,7 @@ from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups
 from aiounifi.interfaces.outlets import Outlets
 from aiounifi.interfaces.port_forwarding import PortForwarding
 from aiounifi.interfaces.ports import Ports
+from aiounifi.interfaces.traffic_routes import TrafficRoutes
 from aiounifi.interfaces.traffic_rules import TrafficRules
 from aiounifi.interfaces.wlans import Wlans
 from aiounifi.models.api import ApiItemT
@@ -31,6 +32,7 @@ from aiounifi.models.event import Event, EventKey
 from aiounifi.models.outlet import Outlet
 from aiounifi.models.port import Port
 from aiounifi.models.port_forward import PortForward, PortForwardEnableRequest
+from aiounifi.models.traffic_route import TrafficRoute, TrafficRouteSaveRequest
 from aiounifi.models.traffic_rule import TrafficRule, TrafficRuleEnableRequest
 from aiounifi.models.wlan import Wlan, WlanEnableRequest
 
@@ -170,6 +172,16 @@ async def async_traffic_rule_control_fn(
     await hub.api.traffic_rules.update()
 
 
+async def async_traffic_route_control_fn(
+    hub: UnifiHub, obj_id: str, target: bool
+) -> None:
+    """Control traffic route state."""
+    traffic_route = hub.api.traffic_routes[obj_id].raw
+    await hub.api.request(TrafficRouteSaveRequest.create(traffic_route, target))
+    # Update the traffic routes so the UI is updated appropriately
+    await hub.api.traffic_routes.update()
+
+
 async def async_wlan_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
     """Control outlet relay."""
     await hub.api.request(WlanEnableRequest.create(obj_id, target))
@@ -263,6 +275,19 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         object_fn=lambda api, obj_id: api.traffic_rules[obj_id],
         unique_id_fn=lambda hub, obj_id: f"traffic_rule-{obj_id}",
     ),
+    UnifiSwitchEntityDescription[TrafficRoutes, TrafficRoute](
+        key="Traffic route control",
+        translation_key="traffic_route_control",
+        device_class=SwitchDeviceClass.SWITCH,
+        entity_category=EntityCategory.CONFIG,
+        api_handler_fn=lambda api: api.traffic_routes,
+        control_fn=async_traffic_route_control_fn,
+        device_info_fn=async_unifi_network_device_info_fn,
+        is_on_fn=lambda hub, traffic_route: traffic_route.enabled,
+        name_fn=lambda traffic_route: traffic_route.description,
+        object_fn=lambda api, obj_id: api.traffic_routes[obj_id],
+        unique_id_fn=lambda hub, obj_id: f"traffic_route-{obj_id}",
+    ),
     UnifiSwitchEntityDescription[Ports, Port](
         key="PoE port control",
         translation_key="poe_port_control",

tests/components/unifi/conftest.py

@@ -174,6 +174,7 @@ def fixture_request(
     dpi_group_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]],
     site_payload: list[dict[str, Any]],
     system_information_payload: list[dict[str, Any]],
     wlan_payload: list[dict[str, Any]],
@@ -214,6 +215,7 @@ def fixture_request(
         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)
         mock_get_request(f"/v2/api/site/{site_id}/trafficrules", traffic_rule_payload)
+        mock_get_request(f"/v2/api/site/{site_id}/trafficroutes", traffic_route_payload)
 
     return __mock_requests
 
@@ -291,6 +293,12 @@ def traffic_rule_payload_data() -> list[dict[str, Any]]:
     return []
 
 
+@pytest.fixture(name="traffic_route_payload")
+def traffic_route_payload_data() -> list[dict[str, Any]]:
+    """Traffic route data."""
+    return []
+
+
 @pytest.fixture(name="wlan_payload")
 def fixture_wlan_data() -> list[dict[str, Any]]:
     """WLAN data."""

tests/components/unifi/test_switch.py

@@ -809,6 +809,24 @@ TRAFFIC_RULE = {
     "target_devices": [{"client_mac": CLIENT_1["mac"], "type": "CLIENT"}],
 }
 
+TRAFFIC_ROUTE = {
+    "_id": "676f8dbb8f1d54503bba19ab",
+    "description": "Test traffic route",
+    "domains": [{"domain": "youtube.com", "port_ranges": [], "ports": []}],
+    "enabled": True,
+    "ip_addresses": [],
+    "ip_ranges": [],
+    "kill_switch_enabled": True,
+    "matching_target": "DOMAIN",
+    "network_id": "676f8d288f1d54503bba1987",
+    "next_hop": "",
+    "regions": [],
+    "target_devices": [
+        {"network_id": "6060b00f45de3905133cea14", "type": "NETWORK"},
+        {"network_id": "6060ae6045de3905133cea0a", "type": "NETWORK"},
+    ],
+}
+
 
 @pytest.mark.parametrize(
     "config_entry_options", [{CONF_BLOCK_CLIENT: [BLOCKED["mac"]]}]
@@ -1154,6 +1172,60 @@ async def test_traffic_rules(
     assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call
 
 
+@pytest.mark.parametrize(("traffic_route_payload"), [([TRAFFIC_ROUTE])])
+async def test_traffic_routes(
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    config_entry_setup: MockConfigEntry,
+    traffic_route_payload: list[dict[str, Any]],
+) -> None:
+    """Test control of UniFi traffic routes."""
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
+
+    # Validate state object
+    assert hass.states.get("switch.unifi_network_test_traffic_route").state == STATE_ON
+
+    traffic_route = deepcopy(traffic_route_payload[0])
+
+    # Disable traffic route
+    aioclient_mock.put(
+        f"https://{config_entry_setup.data[CONF_HOST]}:1234"
+        f"/v2/api/site/{config_entry_setup.data[CONF_SITE_ID]}"
+        f"/trafficroutes/{traffic_route['_id']}",
+    )
+
+    call_count = aioclient_mock.call_count
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        "turn_off",
+        {"entity_id": "switch.unifi_network_test_traffic_route"},
+        blocking=True,
+    )
+    # Updating the value for traffic routes will make another call to retrieve the values
+    assert aioclient_mock.call_count == call_count + 2
+    expected_disable_call = deepcopy(traffic_route)
+    expected_disable_call["enabled"] = False
+
+    assert aioclient_mock.mock_calls[call_count][2] == expected_disable_call
+
+    call_count = aioclient_mock.call_count
+
+    # Enable traffic route
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        "turn_on",
+        {"entity_id": "switch.unifi_network_test_traffic_route"},
+        blocking=True,
+    )
+
+    expected_enable_call = deepcopy(traffic_route)
+    expected_enable_call["enabled"] = True
+
+    assert aioclient_mock.call_count == call_count + 2
+    assert aioclient_mock.mock_calls[call_count][2] == expected_enable_call
+
+
 @pytest.mark.parametrize(
     ("device_payload", "entity_id", "outlet_index", "expected_switches"),
     [