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

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()

homeassistant/components/unifi/const.py

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

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()

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",

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

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

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<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'
+# ---

tests/components/unifi/test_config_flow.py

@@ -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 = [

tests/components/unifi/test_controller.py

@@ -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]

tests/components/unifi/test_image.py

@@ -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