Add WLAN QR code support to UniFi Image platform (#97171)

This commit is contained in:
Robert Svensson 2023-07-25 10:11:48 +02:00 committed by GitHub
parent f2726527f2
commit 06f97679ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 282 additions and 9 deletions

View File

@ -308,7 +308,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
return await self.async_step_client_control()
ssids = (
set(self.controller.api.wlans)
{wlan.name for wlan in self.controller.api.wlans.values()}
| {
f"{wlan.name}{wlan.name_combine_suffix}"
for wlan in self.controller.api.wlans.values()

View File

@ -9,6 +9,7 @@ DOMAIN = "unifi"
PLATFORMS = [
Platform.DEVICE_TRACKER,
Platform.IMAGE,
Platform.SENSOR,
Platform.SWITCH,
Platform.UPDATE,

View File

@ -0,0 +1,136 @@
"""Image platform for UniFi Network integration.
Support for QR code for guest WLANs.
"""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import Generic
import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.wlans import Wlans
from aiounifi.models.api import ApiItemT
from aiounifi.models.wlan import Wlan
from homeassistant.components.image import DOMAIN, ImageEntity, ImageEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import DeviceEntryType
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.util.dt as dt_util
from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
from .controller import UniFiController
from .entity import HandlerT, UnifiEntity, UnifiEntityDescription
@callback
def async_wlan_qr_code_image_fn(controller: UniFiController, wlan: Wlan) -> bytes:
"""Calculate receiving data transfer value."""
return controller.api.wlans.generate_wlan_qr_code(wlan)
@callback
def async_wlan_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo:
"""Create device registry entry for WLAN."""
wlan = api.wlans[obj_id]
return DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, wlan.id)},
manufacturer=ATTR_MANUFACTURER,
model="UniFi Network",
name=wlan.name,
)
@dataclass
class UnifiImageEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
"""Validate and load entities from different UniFi handlers."""
image_fn: Callable[[UniFiController, ApiItemT], bytes]
value_fn: Callable[[ApiItemT], str]
@dataclass
class UnifiImageEntityDescription(
ImageEntityDescription,
UnifiEntityDescription[HandlerT, ApiItemT],
UnifiImageEntityDescriptionMixin[HandlerT, ApiItemT],
):
"""Class describing UniFi image entity."""
ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
UnifiImageEntityDescription[Wlans, Wlan](
key="WLAN QR Code",
entity_category=EntityCategory.DIAGNOSTIC,
has_entity_name=True,
entity_registry_enabled_default=False,
allowed_fn=lambda controller, obj_id: True,
api_handler_fn=lambda api: api.wlans,
available_fn=lambda controller, _: controller.available,
device_info_fn=async_wlan_device_info_fn,
event_is_on=None,
event_to_subscribe=None,
name_fn=lambda _: "QR Code",
object_fn=lambda api, obj_id: api.wlans[obj_id],
supported_fn=lambda controller, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"qr_code-{obj_id}",
image_fn=async_wlan_qr_code_image_fn,
value_fn=lambda obj: obj.x_passphrase,
),
)
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up image platform for UniFi Network integration."""
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.register_platform_add_entities(
UnifiImageEntity, ENTITY_DESCRIPTIONS, async_add_entities
)
class UnifiImageEntity(UnifiEntity[HandlerT, ApiItemT], ImageEntity):
"""Base representation of a UniFi image."""
entity_description: UnifiImageEntityDescription[HandlerT, ApiItemT]
_attr_content_type = "image/png"
current_image: bytes | None = None
previous_value = ""
def __init__(
self,
obj_id: str,
controller: UniFiController,
description: UnifiEntityDescription[HandlerT, ApiItemT],
) -> None:
"""Initiatlize UniFi Image entity."""
super().__init__(obj_id, controller, description)
ImageEntity.__init__(self, controller.hass)
def image(self) -> bytes | None:
"""Return bytes of image."""
if self.current_image is None:
description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id)
self.current_image = description.image_fn(self.controller, obj)
return self.current_image
@callback
def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
"""Update entity state."""
description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id)
if (value := description.value_fn(obj)) != self.previous_value:
self.previous_value = value
self.current_image = None
self._attr_image_last_updated = dt_util.utcnow()

View File

@ -8,7 +8,7 @@
"iot_class": "local_push",
"loggers": ["aiounifi"],
"quality_scale": "platinum",
"requirements": ["aiounifi==49"],
"requirements": ["aiounifi==50"],
"ssdp": [
{
"manufacturer": "Ubiquiti Networks",

View File

@ -357,7 +357,7 @@ aiosyncthing==0.5.1
aiotractive==0.5.5
# homeassistant.components.unifi
aiounifi==49
aiounifi==50
# homeassistant.components.vlc_telnet
aiovlc==0.1.0

View File

@ -332,7 +332,7 @@ aiosyncthing==0.5.1
aiotractive==0.5.5
# homeassistant.components.unifi
aiounifi==49
aiounifi==50
# homeassistant.components.vlc_telnet
aiovlc==0.1.0

View File

@ -0,0 +1,7 @@
# serializer version: 1
# name: test_wlan_qr_code
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x84\x00\x00\x00\x84\x01\x00\x00\x00\x00y?\xbe\n\x00\x00\x00\xcaIDATx\xda\xedV[\n\xc30\x0c\x13\xbb\x80\xef\x7fK\xdd\xc0\x93\x94\xfd\xac\x1fcL\xfbl(\xc4\x04*\xacG\xdcb/\x8b\xb8O\xdeO\x00\xccP\x95\x8b\xe5\x03\xd7\xf5\xcd\x89pF\xcf\x8c \\48\x08\nS\x948\x03p\xfe\x80C\xa8\x9d\x16\xc7P\xabvJ}\xe2\xd7\x84[\xe5W\xfc7\xbbS\xfd\xde\xcfB\xf115\xa2\xe3%\x99\xad\x93\xa0:\xbf6\xbeS\xec\x1a^\xb4\xed\xfb\xb2\xab\xd1\x99\xc9\xcdAjx\x89\x0e\xc5\xea\xf4T\xf9\xee\xe40m58\xb6<\x1b\xab~\xf4\xban\xd7:\xceu\x9e\x05\xc4I\xa6\xbb\xfb%q<7:\xbf\xa2\x90wo\xf5<t_\xf5\\]\xe7\xa7\xe5\x15\x1c\xa5\xd1\xdf\x8cV\x1f#,*\x9c\xccC+S\xce\x9f\xfb\xaf\xe0\xc3\xc9\x13/\xb7\x08A\x10\xbe\x06\xd0\x00\x00\x00\x00IEND\xaeB`\x82'
# ---
# name: test_wlan_qr_code.1
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x94\x00\x00\x00\x94\x01\x00\x00\x00\x00]G=y\x00\x00\x00\xfdIDATx\xda\xedV1\x8e\x041\x0cB\xf7\x01\xff\xff\x97\xfc\xc0\x0bd\xb6\xda\xe6\xeeB\xb9V\xa4dR \xc7`<\xd8\x8f \xbew\x7f\xb9\x030\x98!\xb5\xe9\xb8\xfc\xc1g\xfc\xf6Nx\xa3%\x9c\x84\xbf\xae\xf1\x84\xb5 \xe796\xf0\\\npjx~1[xZ\\\xbfy+\xf5\xc3\x9b\x8c\xe9\xf0\xeb\xd0k]\xbe\xa3\xa1\xeb\xfaI\x850\xa2Ex\x9f\x1f-\xeb\xe46!\xba\xc0G\x18\xde\xb0|\x8f\x07e8\xca\xd0\xc0,\xd4/\xed&PA\x1a\xf5\xbe~R2m\x07\x8fa\\\xe3\x9d\xc4DnG\x7f\xb0F&\xc4L\xa3~J\xcciy\xdfF\xff\x9a`i\xda$w\xfcom\xcc\x02Kw\x14\xf4\xc2\xd3fn\xba-\xf0A&A\xe2\x0c\x92\x8e\xbfL<\xcb.\xd8\xf1?0~o\xc14\xfcy\xdc\xc48\xa6\xd0\x98\x1f\x99\xbd\xfb\xd0\xd3\x98o\xd1tFR\x07\x8f\xe95lo\xbeE\x88`\x8f\xdf\x8c`lE\x7f\xdf\xff\xc4\x7f\xde\xbd\x00\xfc\xb3\x80\x95k\x06#\x19\x00\x00\x00\x00IEND\xaeB`\x82'
# ---

View File

@ -74,9 +74,14 @@ DEVICES = [
]
WLANS = [
{"name": "SSID 1"},
{"name": "SSID 2", "name_combine_enabled": False, "name_combine_suffix": "_IOT"},
{"name": "SSID 4", "name_combine_enabled": False},
{"_id": "1", "name": "SSID 1"},
{
"_id": "2",
"name": "SSID 2",
"name_combine_enabled": False,
"name_combine_suffix": "_IOT",
},
{"_id": "3", "name": "SSID 4", "name_combine_enabled": False},
]
DPI_GROUPS = [

View File

@ -10,6 +10,7 @@ from aiounifi.websocket import WebsocketState
import pytest
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.image import DOMAIN as IMAGE_DOMAIN
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.components.unifi.const import (
@ -222,8 +223,9 @@ async def test_controller_setup(
entry = controller.config_entry
assert len(forward_entry_setup.mock_calls) == len(PLATFORMS)
assert forward_entry_setup.mock_calls[0][1] == (entry, TRACKER_DOMAIN)
assert forward_entry_setup.mock_calls[1][1] == (entry, SENSOR_DOMAIN)
assert forward_entry_setup.mock_calls[2][1] == (entry, SWITCH_DOMAIN)
assert forward_entry_setup.mock_calls[1][1] == (entry, IMAGE_DOMAIN)
assert forward_entry_setup.mock_calls[2][1] == (entry, SENSOR_DOMAIN)
assert forward_entry_setup.mock_calls[3][1] == (entry, SWITCH_DOMAIN)
assert controller.host == ENTRY_CONFIG[CONF_HOST]
assert controller.site == ENTRY_CONFIG[CONF_SITE_ID]

View File

@ -0,0 +1,122 @@
"""UniFi Network image platform tests."""
from copy import deepcopy
from datetime import timedelta
from http import HTTPStatus
from aiounifi.models.message import MessageKey
from syrupy.assertion import SnapshotAssertion
from homeassistant.components.image import DOMAIN as IMAGE_DOMAIN
from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
from homeassistant.const import (
EntityCategory,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity_registry import RegistryEntryDisabler
from homeassistant.util import dt as dt_util
from .test_controller import (
setup_unifi_integration,
)
from tests.common import async_fire_time_changed
from tests.test_util.aiohttp import AiohttpClientMocker
from tests.typing import ClientSessionGenerator
WLAN = {
"_id": "012345678910111213141516",
"bc_filter_enabled": False,
"bc_filter_list": [],
"dtim_mode": "default",
"dtim_na": 1,
"dtim_ng": 1,
"enabled": True,
"group_rekey": 3600,
"mac_filter_enabled": False,
"mac_filter_list": [],
"mac_filter_policy": "allow",
"minrate_na_advertising_rates": False,
"minrate_na_beacon_rate_kbps": 6000,
"minrate_na_data_rate_kbps": 6000,
"minrate_na_enabled": False,
"minrate_na_mgmt_rate_kbps": 6000,
"minrate_ng_advertising_rates": False,
"minrate_ng_beacon_rate_kbps": 1000,
"minrate_ng_data_rate_kbps": 1000,
"minrate_ng_enabled": False,
"minrate_ng_mgmt_rate_kbps": 1000,
"name": "SSID 1",
"no2ghz_oui": False,
"schedule": [],
"security": "wpapsk",
"site_id": "5a32aa4ee4b0412345678910",
"usergroup_id": "012345678910111213141518",
"wep_idx": 1,
"wlangroup_id": "012345678910111213141519",
"wpa_enc": "ccmp",
"wpa_mode": "wpa2",
"x_iapp_key": "01234567891011121314151617181920",
"x_passphrase": "password",
}
async def test_wlan_qr_code(
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
hass_client: ClientSessionGenerator,
snapshot: SnapshotAssertion,
mock_unifi_websocket,
) -> None:
"""Test the update_clients function when no clients are found."""
await setup_unifi_integration(hass, aioclient_mock, wlans_response=[WLAN])
assert len(hass.states.async_entity_ids(IMAGE_DOMAIN)) == 0
ent_reg = er.async_get(hass)
ent_reg_entry = ent_reg.async_get("image.ssid_1_qr_code")
assert ent_reg_entry.unique_id == "qr_code-012345678910111213141516"
assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC
# Enable entity
ent_reg.async_update_entity(entity_id="image.ssid_1_qr_code", disabled_by=None)
await hass.async_block_till_done()
async_fire_time_changed(
hass,
dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1),
)
await hass.async_block_till_done()
# Validate state object
image_state_1 = hass.states.get("image.ssid_1_qr_code")
assert image_state_1.name == "SSID 1 QR Code"
# Validate image
client = await hass_client()
resp = await client.get("/api/image_proxy/image.ssid_1_qr_code")
assert resp.status == HTTPStatus.OK
body = await resp.read()
assert body == snapshot
# Update state object - same password - no change to state
mock_unifi_websocket(message=MessageKey.WLAN_CONF_UPDATED, data=WLAN)
await hass.async_block_till_done()
image_state_2 = hass.states.get("image.ssid_1_qr_code")
assert image_state_1.state == image_state_2.state
# Update state object - changeed password - new state
data = deepcopy(WLAN)
data["x_passphrase"] = "new password"
mock_unifi_websocket(message=MessageKey.WLAN_CONF_UPDATED, data=data)
await hass.async_block_till_done()
image_state_3 = hass.states.get("image.ssid_1_qr_code")
assert image_state_1.state != image_state_3.state
# Validate image
client = await hass_client()
resp = await client.get("/api/image_proxy/image.ssid_1_qr_code")
assert resp.status == HTTPStatus.OK
body = await resp.read()
assert body == snapshot