From 06f97679ee8b7f675ae53dbd9a7990e1a2993501 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 25 Jul 2023 10:11:48 +0200 Subject: [PATCH] Add WLAN QR code support to UniFi Image platform (#97171) --- homeassistant/components/unifi/config_flow.py | 2 +- homeassistant/components/unifi/const.py | 1 + homeassistant/components/unifi/image.py | 136 ++++++++++++++++++ homeassistant/components/unifi/manifest.json | 2 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- .../unifi/snapshots/test_image.ambr | 7 + tests/components/unifi/test_config_flow.py | 11 +- tests/components/unifi/test_controller.py | 6 +- tests/components/unifi/test_image.py | 122 ++++++++++++++++ 10 files changed, 282 insertions(+), 9 deletions(-) create mode 100644 homeassistant/components/unifi/image.py create mode 100644 tests/components/unifi/snapshots/test_image.ambr create mode 100644 tests/components/unifi/test_image.py diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index d283b668995..12f2d49e416 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -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() diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index b5cea06c719..e03bd50d483 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -9,6 +9,7 @@ DOMAIN = "unifi" PLATFORMS = [ Platform.DEVICE_TRACKER, + Platform.IMAGE, Platform.SENSOR, Platform.SWITCH, Platform.UPDATE, diff --git a/homeassistant/components/unifi/image.py b/homeassistant/components/unifi/image.py new file mode 100644 index 00000000000..730720753d4 --- /dev/null +++ b/homeassistant/components/unifi/image.py @@ -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() diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index 9bfb01e5a88..c34d1035158 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -8,7 +8,7 @@ "iot_class": "local_push", "loggers": ["aiounifi"], "quality_scale": "platinum", - "requirements": ["aiounifi==49"], + "requirements": ["aiounifi==50"], "ssdp": [ { "manufacturer": "Ubiquiti Networks", diff --git a/requirements_all.txt b/requirements_all.txt index 731c4f18c45..cde697360aa 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 426f249c2f3..70e07baf8d6 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -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 diff --git a/tests/components/unifi/snapshots/test_image.ambr b/tests/components/unifi/snapshots/test_image.ambr new file mode 100644 index 00000000000..77b171118a1 --- /dev/null +++ b/tests/components/unifi/snapshots/test_image.ambr @@ -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 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