diff --git a/homeassistant/components/unifi/hub/entity_loader.py b/homeassistant/components/unifi/hub/entity_loader.py index f11ddefec98..64403152b0c 100644 --- a/homeassistant/components/unifi/hub/entity_loader.py +++ b/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( diff --git a/homeassistant/components/unifi/icons.json b/homeassistant/components/unifi/icons.json index 76990c1c4a1..6874bb5ae03 100644 --- a/homeassistant/components/unifi/icons.json +++ b/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": { diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 01843a8a95b..7741e57c82c 100644 --- a/homeassistant/components/unifi/switch.py +++ b/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", diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index 798b613b18d..702f8629219 100644 --- a/tests/components/unifi/conftest.py +++ b/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.""" diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index cb5dcdac428..e4765d1181e 100644 --- a/tests/components/unifi/test_switch.py +++ b/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"), [