From 53131390ac7640a83c1bea0e9b8acc21b5e90eb8 Mon Sep 17 00:00:00 2001
From: Robert Svensson <Kane610@users.noreply.github.com>
Date: Fri, 26 Jul 2024 19:22:09 +0200
Subject: [PATCH] Use snapshot in UniFi image tests (#122608)

* Use snapshot in UniFi image tests

* Make Image access_token deterministic
---
 .../unifi/snapshots/test_image.ambr           | 48 ++++++++++++
 tests/components/unifi/test_image.py          | 77 +++++++++++++++----
 2 files changed, 108 insertions(+), 17 deletions(-)

diff --git a/tests/components/unifi/snapshots/test_image.ambr b/tests/components/unifi/snapshots/test_image.ambr
index 83d76688ea3..e33ec678217 100644
--- a/tests/components/unifi/snapshots/test_image.ambr
+++ b/tests/components/unifi/snapshots/test_image.ambr
@@ -1,4 +1,52 @@
 # serializer version: 1
+# name: test_entity_and_device_data[site_payload0-wlan_payload0][image.ssid_1_qr_code-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'image',
+    'entity_category': <EntityCategory.DIAGNOSTIC: 'diagnostic'>,
+    'entity_id': 'image.ssid_1_qr_code',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'QR Code',
+    'platform': 'unifi',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': None,
+    'unique_id': 'qr_code-012345678910111213141516',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_entity_and_device_data[site_payload0-wlan_payload0][image.ssid_1_qr_code-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'access_token': '1',
+      'entity_picture': '/api/image_proxy/image.ssid_1_qr_code?token=1',
+      'friendly_name': 'SSID 1 QR Code',
+    }),
+    'context': <ANY>,
+    'entity_id': 'image.ssid_1_qr_code',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2021-01-01T01:01:00+00:00',
+  })
+# ---
 # 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'
 # ---
diff --git a/tests/components/unifi/test_image.py b/tests/components/unifi/test_image.py
index 75d2f02900d..6733845c52f 100644
--- a/tests/components/unifi/test_image.py
+++ b/tests/components/unifi/test_image.py
@@ -3,22 +3,35 @@
 from copy import deepcopy
 from datetime import timedelta
 from http import HTTPStatus
+from typing import Any
+from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
 import pytest
-from syrupy.assertion import SnapshotAssertion
+from syrupy import SnapshotAssertion
 
 from homeassistant.components.image import DOMAIN as IMAGE_DOMAIN
 from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
-from homeassistant.const import STATE_UNAVAILABLE, EntityCategory
+from homeassistant.const import STATE_UNAVAILABLE, Platform
 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 tests.common import async_fire_time_changed
+from tests.common import async_fire_time_changed, snapshot_platform
 from tests.typing import ClientSessionGenerator
 
+
+@pytest.fixture(autouse=True)
+def mock_getrandbits():
+    """Mock image access token which normally is randomized."""
+    with patch(
+        "homeassistant.components.image.SystemRandom.getrandbits",
+        return_value=1,
+    ):
+        yield
+
+
 WLAN = {
     "_id": "012345678910111213141516",
     "bc_filter_enabled": False,
@@ -56,6 +69,32 @@ WLAN = {
 }
 
 
+@pytest.mark.parametrize("wlan_payload", [[WLAN]])
+@pytest.mark.parametrize(
+    "site_payload",
+    [
+        [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}],
+        [{"desc": "Site name", "name": "site_id", "role": "not admin", "_id": "1"}],
+    ],
+)
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+@pytest.mark.freeze_time("2021-01-01 01:01:00")
+async def test_entity_and_device_data(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    config_entry_factory,
+    site_payload: dict[str, Any],
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Validate entity and device data with and without admin rights."""
+    with patch("homeassistant.components.unifi.PLATFORMS", [Platform.IMAGE]):
+        config_entry = await config_entry_factory()
+    if site_payload[0]["role"] == "admin":
+        await snapshot_platform(hass, entity_registry, snapshot, config_entry.entry_id)
+    else:
+        assert len(hass.states.async_entity_ids(IMAGE_DOMAIN)) == 0
+
+
 @pytest.mark.parametrize("wlan_payload", [[WLAN]])
 @pytest.mark.usefixtures("config_entry_setup")
 async def test_wlan_qr_code(
@@ -64,15 +103,12 @@ async def test_wlan_qr_code(
     hass_client: ClientSessionGenerator,
     snapshot: SnapshotAssertion,
     mock_websocket_message,
-    mock_websocket_state,
 ) -> None:
     """Test the update_clients function when no clients are found."""
     assert len(hass.states.async_entity_ids(IMAGE_DOMAIN)) == 0
 
     ent_reg_entry = entity_registry.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
     entity_registry.async_update_entity(
@@ -84,10 +120,6 @@ async def test_wlan_qr_code(
     )
     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")
@@ -96,8 +128,8 @@ async def test_wlan_qr_code(
     assert body == snapshot
 
     # Update state object - same password - no change to state
+    image_state_1 = hass.states.get("image.ssid_1_qr_code")
     mock_websocket_message(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
 
@@ -105,7 +137,6 @@ async def test_wlan_qr_code(
     data = deepcopy(WLAN)
     data["x_passphrase"] = "new password"
     mock_websocket_message(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
 
@@ -116,25 +147,37 @@ async def test_wlan_qr_code(
     body = await resp.read()
     assert body == snapshot
 
-    # Availability signalling
 
-    # Controller disconnects
+@pytest.mark.parametrize("wlan_payload", [[WLAN]])
+@pytest.mark.usefixtures("config_entry_setup")
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> None:
+    """Verify entities state reflect on hub becoming unavailable."""
+    assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
+
+    # Controller unavailable
     await mock_websocket_state.disconnect()
     assert hass.states.get("image.ssid_1_qr_code").state == STATE_UNAVAILABLE
 
-    # Controller reconnects
+    # Controller available
     await mock_websocket_state.reconnect()
     assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
 
+
+@pytest.mark.parametrize("wlan_payload", [[WLAN]])
+@pytest.mark.usefixtures("config_entry_setup")
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_source_availability(hass: HomeAssistant, mock_websocket_message) -> None:
+    """Verify entities state reflect on source becoming unavailable."""
+    assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
+
     # WLAN gets disabled
     wlan_1 = deepcopy(WLAN)
     wlan_1["enabled"] = False
     mock_websocket_message(message=MessageKey.WLAN_CONF_UPDATED, data=wlan_1)
-    await hass.async_block_till_done()
     assert hass.states.get("image.ssid_1_qr_code").state == STATE_UNAVAILABLE
 
     # WLAN gets re-enabled
     wlan_1["enabled"] = True
     mock_websocket_message(message=MessageKey.WLAN_CONF_UPDATED, data=wlan_1)
-    await hass.async_block_till_done()
     assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE