Unifi add port forward control to switch platform (#98309)

This commit is contained in:
Robert Svensson 2023-08-21 22:01:44 +02:00 committed by GitHub
parent 78f0d8bc9c
commit d0d160f11c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 144 additions and 9 deletions

View File

@ -53,12 +53,12 @@ def async_wlan_available_fn(controller: UniFiController, obj_id: str) -> bool:
@callback @callback
def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: def async_device_device_info_fn(controller: UniFiController, obj_id: str) -> DeviceInfo:
"""Create device registry entry for device.""" """Create device registry entry for device."""
if "_" in obj_id: # Sub device (outlet or port) if "_" in obj_id: # Sub device (outlet or port)
obj_id = obj_id.partition("_")[0] obj_id = obj_id.partition("_")[0]
device = api.devices[obj_id] device = controller.api.devices[obj_id]
return DeviceInfo( return DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, device.mac)}, connections={(CONNECTION_NETWORK_MAC, device.mac)},
manufacturer=ATTR_MANUFACTURER, manufacturer=ATTR_MANUFACTURER,
@ -70,9 +70,9 @@ def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> Device
@callback @callback
def async_wlan_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: def async_wlan_device_info_fn(controller: UniFiController, obj_id: str) -> DeviceInfo:
"""Create device registry entry for WLAN.""" """Create device registry entry for WLAN."""
wlan = api.wlans[obj_id] wlan = controller.api.wlans[obj_id]
return DeviceInfo( return DeviceInfo(
entry_type=DeviceEntryType.SERVICE, entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, wlan.id)}, identifiers={(DOMAIN, wlan.id)},
@ -83,9 +83,9 @@ def async_wlan_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceIn
@callback @callback
def async_client_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: def async_client_device_info_fn(controller: UniFiController, obj_id: str) -> DeviceInfo:
"""Create device registry entry for client.""" """Create device registry entry for client."""
client = api.clients[obj_id] client = controller.api.clients[obj_id]
return DeviceInfo( return DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, obj_id)}, connections={(CONNECTION_NETWORK_MAC, obj_id)},
default_manufacturer=client.oui, default_manufacturer=client.oui,
@ -100,7 +100,7 @@ class UnifiDescription(Generic[HandlerT, ApiItemT]):
allowed_fn: Callable[[UniFiController, str], bool] allowed_fn: Callable[[UniFiController, str], bool]
api_handler_fn: Callable[[aiounifi.Controller], HandlerT] api_handler_fn: Callable[[aiounifi.Controller], HandlerT]
available_fn: Callable[[UniFiController, str], bool] available_fn: Callable[[UniFiController, str], bool]
device_info_fn: Callable[[aiounifi.Controller, str], DeviceInfo | None] device_info_fn: Callable[[UniFiController, str], DeviceInfo | None]
event_is_on: tuple[EventKey, ...] | None event_is_on: tuple[EventKey, ...] | None
event_to_subscribe: tuple[EventKey, ...] | None event_to_subscribe: tuple[EventKey, ...] | None
name_fn: Callable[[ApiItemT], str | None] name_fn: Callable[[ApiItemT], str | None]
@ -137,7 +137,7 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
self._removed = False self._removed = False
self._attr_available = description.available_fn(controller, obj_id) self._attr_available = description.available_fn(controller, obj_id)
self._attr_device_info = description.device_info_fn(controller.api, obj_id) self._attr_device_info = description.device_info_fn(controller, obj_id)
self._attr_should_poll = description.should_poll self._attr_should_poll = description.should_poll
self._attr_unique_id = description.unique_id_fn(controller, obj_id) self._attr_unique_id = description.unique_id_fn(controller, obj_id)

View File

@ -17,6 +17,7 @@ from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.clients import Clients from aiounifi.interfaces.clients import Clients
from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups
from aiounifi.interfaces.outlets import Outlets from aiounifi.interfaces.outlets import Outlets
from aiounifi.interfaces.port_forwarding import PortForwarding
from aiounifi.interfaces.ports import Ports from aiounifi.interfaces.ports import Ports
from aiounifi.interfaces.wlans import Wlans from aiounifi.interfaces.wlans import Wlans
from aiounifi.models.api import ApiItemT from aiounifi.models.api import ApiItemT
@ -30,6 +31,7 @@ from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup
from aiounifi.models.event import Event, EventKey from aiounifi.models.event import Event, EventKey
from aiounifi.models.outlet import Outlet from aiounifi.models.outlet import Outlet
from aiounifi.models.port import Port from aiounifi.models.port import Port
from aiounifi.models.port_forward import PortForward, PortForwardEnableRequest
from aiounifi.models.wlan import Wlan, WlanEnableRequest from aiounifi.models.wlan import Wlan, WlanEnableRequest
from homeassistant.components.switch import ( from homeassistant.components.switch import (
@ -75,7 +77,9 @@ def async_dpi_group_is_on_fn(
@callback @callback
def async_dpi_group_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: def async_dpi_group_device_info_fn(
controller: UniFiController, obj_id: str
) -> DeviceInfo:
"""Create device registry entry for DPI group.""" """Create device registry entry for DPI group."""
return DeviceInfo( return DeviceInfo(
entry_type=DeviceEntryType.SERVICE, entry_type=DeviceEntryType.SERVICE,
@ -86,6 +90,22 @@ def async_dpi_group_device_info_fn(api: aiounifi.Controller, obj_id: str) -> Dev
) )
@callback
def async_port_forward_device_info_fn(
controller: UniFiController, obj_id: str
) -> DeviceInfo:
"""Create device registry entry for port forward."""
unique_id = controller.config_entry.unique_id
assert unique_id is not None
return DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, unique_id)},
manufacturer=ATTR_MANUFACTURER,
model="UniFi Network",
name="UniFi Network",
)
async def async_block_client_control_fn( async def async_block_client_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool api: aiounifi.Controller, obj_id: str, target: bool
) -> None: ) -> None:
@ -136,6 +156,14 @@ async def async_poe_port_control_fn(
await api.request(DeviceSetPoePortModeRequest.create(device, int(index), state)) await api.request(DeviceSetPoePortModeRequest.create(device, int(index), state))
async def async_port_forward_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool
) -> None:
"""Control port forward state."""
port_forward = api.port_forwarding[obj_id]
await api.request(PortForwardEnableRequest.create(port_forward, target))
async def async_wlan_control_fn( async def async_wlan_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool api: aiounifi.Controller, obj_id: str, target: bool
) -> None: ) -> None:
@ -222,6 +250,26 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
supported_fn=async_outlet_supports_switching_fn, supported_fn=async_outlet_supports_switching_fn,
unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-outlet-{obj_id.split('_', 1)[1]}", unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-outlet-{obj_id.split('_', 1)[1]}",
), ),
UnifiSwitchEntityDescription[PortForwarding, PortForward](
key="Port forward control",
device_class=SwitchDeviceClass.SWITCH,
entity_category=EntityCategory.CONFIG,
has_entity_name=True,
icon="mdi:upload-network",
allowed_fn=lambda controller, obj_id: True,
api_handler_fn=lambda api: api.port_forwarding,
available_fn=lambda controller, obj_id: controller.available,
control_fn=async_port_forward_control_fn,
device_info_fn=async_port_forward_device_info_fn,
event_is_on=None,
event_to_subscribe=None,
is_on_fn=lambda controller, port_forward: port_forward.enabled,
name_fn=lambda port_forward: f"{port_forward.name}",
object_fn=lambda api, obj_id: api.port_forwarding[obj_id],
should_poll=False,
supported_fn=lambda controller, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"port_forward-{obj_id}",
),
UnifiSwitchEntityDescription[Ports, Port]( UnifiSwitchEntityDescription[Ports, Port](
key="PoE port control", key="PoE port control",
device_class=SwitchDeviceClass.OUTLET, device_class=SwitchDeviceClass.OUTLET,

View File

@ -1518,3 +1518,90 @@ async def test_wlan_switches(
mock_unifi_websocket(state=WebsocketState.RUNNING) mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done() await hass.async_block_till_done()
assert hass.states.get("switch.ssid_1").state == STATE_OFF assert hass.states.get("switch.ssid_1").state == STATE_OFF
async def test_port_forwarding_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
) -> None:
"""Test control of UniFi port forwarding."""
_data = {
"_id": "5a32aa4ee4b0412345678911",
"dst_port": "12345",
"enabled": True,
"fwd_port": "23456",
"fwd": "10.0.0.2",
"name": "plex",
"pfwd_interface": "wan",
"proto": "tcp_udp",
"site_id": "5a32aa4ee4b0412345678910",
"src": "any",
}
config_entry = await setup_unifi_integration(
hass, aioclient_mock, port_forward_response=[_data.copy()]
)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("switch.unifi_network_plex")
assert ent_reg_entry.unique_id == "port_forward-5a32aa4ee4b0412345678911"
assert ent_reg_entry.entity_category is EntityCategory.CONFIG
# Validate state object
switch_1 = hass.states.get("switch.unifi_network_plex")
assert switch_1 is not None
assert switch_1.state == STATE_ON
assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.SWITCH
# Update state object
data = _data.copy()
data["enabled"] = False
mock_unifi_websocket(message=MessageKey.PORT_FORWARD_UPDATED, data=data)
await hass.async_block_till_done()
assert hass.states.get("switch.unifi_network_plex").state == STATE_OFF
# Disable port forward
aioclient_mock.clear_requests()
aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}"
+ f"/rest/portforward/{data['_id']}",
)
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_off",
{"entity_id": "switch.unifi_network_plex"},
blocking=True,
)
assert aioclient_mock.call_count == 1
data = _data.copy()
data["enabled"] = False
assert aioclient_mock.mock_calls[0][2] == data
# Enable port forward
await hass.services.async_call(
SWITCH_DOMAIN,
"turn_on",
{"entity_id": "switch.unifi_network_plex"},
blocking=True,
)
assert aioclient_mock.call_count == 2
assert aioclient_mock.mock_calls[1][2] == _data
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
assert hass.states.get("switch.unifi_network_plex").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
assert hass.states.get("switch.unifi_network_plex").state == STATE_OFF
# Remove entity on deleted message
mock_unifi_websocket(message=MessageKey.PORT_FORWARD_DELETED, data=_data)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0