-
-
Notifications
You must be signed in to change notification settings - Fork 28.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add WLAN QR code support to UniFi Image platform (#97171)
- Loading branch information
Showing
10 changed files
with
282 additions
and
9 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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' | ||
# --- |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |