Use snapshot in UniFi switch tests (#122871)

* Use snapshot in UniFi switch tests

* Fix review comment

tests/components/unifi/test_switch.py

@@ -3,15 +3,16 @@
 from copy import deepcopy
 from datetime import timedelta
 from typing import Any
+from unittest.mock import patch
 
 from aiounifi.models.message import MessageKey
 import pytest
+from syrupy import SnapshotAssertion
 
 from homeassistant.components.switch import (
     DOMAIN as SWITCH_DOMAIN,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
-    SwitchDeviceClass,
 )
 from homeassistant.components.unifi.const import (
     CONF_BLOCK_CLIENT,
@@ -23,13 +24,12 @@ from homeassistant.components.unifi.const import (
 )
 from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
 from homeassistant.const import (
-    ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     CONF_HOST,
     STATE_OFF,
     STATE_ON,
     STATE_UNAVAILABLE,
-    EntityCategory,
+    Platform,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
@@ -43,7 +43,7 @@ from .conftest import (
     WebsocketStateManager,
 )
 
-from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 CLIENT_1 = {
@@ -810,6 +810,34 @@ TRAFFIC_RULE = {
 }
 
 
+@pytest.mark.parametrize(
+    "config_entry_options", [{CONF_BLOCK_CLIENT: [BLOCKED["mac"]]}]
+)
+@pytest.mark.parametrize("client_payload", [[BLOCKED]])
+@pytest.mark.parametrize("device_payload", [[DEVICE_1, OUTLET_UP1, PDU_DEVICE_1]])
+@pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
+@pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
+@pytest.mark.parametrize("port_forward_payload", [[PORT_FORWARD_PLEX]])
+@pytest.mark.parametrize(("traffic_rule_payload"), [([TRAFFIC_RULE])])
+@pytest.mark.parametrize("wlan_payload", [[WLAN]])
+@pytest.mark.parametrize(
+    "site_payload",
+    [[{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}]],
+)
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_entity_and_device_data(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    config_entry_factory: ConfigEntryFactoryType,
+    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.SWITCH]):
+        config_entry = await config_entry_factory()
+    await snapshot_platform(hass, entity_registry, snapshot, config_entry.entry_id)
+
+
 @pytest.mark.parametrize("client_payload", [[CONTROLLER_HOST]])
 @pytest.mark.parametrize("device_payload", [[DEVICE_1]])
 @pytest.mark.usefixtures("config_entry_setup")
@@ -819,18 +847,6 @@ async def test_hub_not_client(hass: HomeAssistant) -> None:
     assert hass.states.get("switch.cloud_key") is None
 
 
-@pytest.mark.parametrize("client_payload", [[CLIENT_1]])
-@pytest.mark.parametrize("device_payload", [[DEVICE_1]])
-@pytest.mark.parametrize(
-    "site_payload",
-    [[{"desc": "Site name", "name": "site_id", "role": "not admin", "_id": "1"}]],
-)
-@pytest.mark.usefixtures("config_entry_setup")
-async def test_not_admin(hass: HomeAssistant) -> None:
-    """Test that switch platform only work on an admin account."""
-    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
-
-
 @pytest.mark.parametrize(
     "config_entry_options",
     [
@@ -841,40 +857,17 @@ async def test_not_admin(hass: HomeAssistant) -> None:
         }
     ],
 )
-@pytest.mark.parametrize("client_payload", [[CLIENT_4]])
 @pytest.mark.parametrize("clients_all_payload", [[BLOCKED, UNBLOCKED, CLIENT_1]])
 @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 async def test_switches(
     hass: HomeAssistant,
-    entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
     config_entry_setup: MockConfigEntry,
 ) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
 
-    switch_4 = hass.states.get("switch.poe_client_4")
-    assert switch_4 is None
-
-    blocked = hass.states.get("switch.block_client_1")
-    assert blocked is not None
-    assert blocked.state == "off"
-
-    unblocked = hass.states.get("switch.block_client_2")
-    assert unblocked is not None
-    assert unblocked.state == "on"
-
-    dpi_switch = hass.states.get("switch.block_media_streaming")
-    assert dpi_switch is not None
-    assert dpi_switch.state == "on"
-    assert dpi_switch.attributes["icon"] == "mdi:network"
-
-    for entry_id in ("switch.block_client_1", "switch.block_media_streaming"):
-        assert (
-            entity_registry.async_get(entry_id).entity_category is EntityCategory.CONFIG
-        )
-
     # Block and unblock client
     aioclient_mock.clear_requests()
     aioclient_mock.post(
@@ -1038,10 +1031,7 @@ async def test_dpi_switches(
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
-    dpi_switch = hass.states.get("switch.block_media_streaming")
+    assert hass.states.get("switch.block_media_streaming").state == STATE_ON
-    assert dpi_switch is not None
-    assert dpi_switch.state == STATE_ON
-    assert dpi_switch.attributes["icon"] == "mdi:network"
 
     mock_websocket_message(data=DPI_APP_DISABLED_EVENT)
     await hass.async_block_till_done()
@@ -1118,20 +1108,18 @@ async def test_traffic_rules(
     traffic_rule_payload: list[dict[str, Any]],
 ) -> None:
     """Test control of UniFi traffic rules."""
-
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
     # Validate state object
-    switch_1 = hass.states.get("switch.unifi_network_test_traffic_rule")
+    assert hass.states.get("switch.unifi_network_test_traffic_rule").state == STATE_ON
-    assert switch_1.state == STATE_ON
-    assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.SWITCH
 
     traffic_rule = deepcopy(traffic_rule_payload[0])
 
     # Disable traffic rule
     aioclient_mock.put(
         f"https://{config_entry_setup.data[CONF_HOST]}:1234"
-        f"/v2/api/site/{config_entry_setup.data[CONF_SITE_ID]}/trafficrules/{traffic_rule['_id']}",
+        f"/v2/api/site/{config_entry_setup.data[CONF_SITE_ID]}"
+        f"/trafficrules/{traffic_rule['_id']}",
     )
 
     call_count = aioclient_mock.call_count
@@ -1188,10 +1176,7 @@ async def test_outlet_switches(
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == expected_switches
 
     # Validate state object
-    switch_1 = hass.states.get(f"switch.{entity_id}")
+    assert hass.states.get(f"switch.{entity_id}").state == STATE_ON
-    assert switch_1 is not None
-    assert switch_1.state == STATE_ON
-    assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.OUTLET
 
     # Update state object
     device_1 = deepcopy(device_payload[0])
@@ -1250,15 +1235,6 @@ async def test_outlet_switches(
     await hass.async_block_till_done()
     assert hass.states.get(f"switch.{entity_id}").state == STATE_OFF
 
-    # Unload config entry
-    await hass.config_entries.async_unload(config_entry_setup.entry_id)
-    assert hass.states.get(f"switch.{entity_id}").state == STATE_UNAVAILABLE
-
-    # Remove config entry
-    await hass.config_entries.async_remove(config_entry_setup.entry_id)
-    await hass.async_block_till_done()
-    assert hass.states.get(f"switch.{entity_id}") is None
-
 
 @pytest.mark.parametrize(
     "config_entry_options",
@@ -1359,8 +1335,8 @@ async def test_poe_port_switches(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message: WebsocketMessageMock,
     config_entry_setup: MockConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Test PoE port entities work."""
@@ -1368,7 +1344,6 @@ async def test_poe_port_switches(
 
     ent_reg_entry = entity_registry.async_get("switch.mock_name_port_1_poe")
     assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION
-    assert ent_reg_entry.entity_category is EntityCategory.CONFIG
 
     # Enable entity
     entity_registry.async_update_entity(
@@ -1385,10 +1360,7 @@ async def test_poe_port_switches(
     await hass.async_block_till_done()
 
     # Validate state object
-    switch_1 = hass.states.get("switch.mock_name_port_1_poe")
+    assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_ON
-    assert switch_1 is not None
-    assert switch_1.state == STATE_ON
-    assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.OUTLET
 
     # Update state object
     device_1 = deepcopy(device_payload[0])
@@ -1456,24 +1428,16 @@ async def test_poe_port_switches(
 @pytest.mark.parametrize("wlan_payload", [[WLAN]])
 async def test_wlan_switches(
     hass: HomeAssistant,
-    entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message: WebsocketMessageMock,
     config_entry_setup: MockConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
     wlan_payload: list[dict[str, Any]],
 ) -> None:
     """Test control of UniFi WLAN availability."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
-    ent_reg_entry = entity_registry.async_get("switch.ssid_1")
-    assert ent_reg_entry.unique_id == "wlan-012345678910111213141516"
-    assert ent_reg_entry.entity_category is EntityCategory.CONFIG
-
     # Validate state object
-    switch_1 = hass.states.get("switch.ssid_1")
+    assert hass.states.get("switch.ssid_1").state == STATE_ON
-    assert switch_1 is not None
-    assert switch_1.state == STATE_ON
-    assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.SWITCH
 
     # Update state object
     wlan = deepcopy(wlan_payload[0])
@@ -1512,24 +1476,16 @@ async def test_wlan_switches(
 @pytest.mark.parametrize("port_forward_payload", [[PORT_FORWARD_PLEX]])
 async def test_port_forwarding_switches(
     hass: HomeAssistant,
-    entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message: WebsocketMessageMock,
     config_entry_setup: MockConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
     port_forward_payload: list[dict[str, Any]],
 ) -> None:
     """Test control of UniFi port forwarding."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
-    ent_reg_entry = entity_registry.async_get("switch.unifi_network_plex")
-    assert ent_reg_entry.unique_id == "port_forward-5a32aa4ee4b0412345678911"
-    assert ent_reg_entry.entity_category is EntityCategory.CONFIG
-
     # Validate state object
-    switch_1 = hass.states.get("switch.unifi_network_plex")
+    assert hass.states.get("switch.unifi_network_plex").state == STATE_ON
-    assert switch_1 is not None
-    assert switch_1.state == STATE_ON
-    assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.SWITCH
 
     # Update state object
     data = port_forward_payload[0].copy()
@@ -1648,6 +1604,7 @@ async def test_updating_unique_id(
 @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 @pytest.mark.parametrize("port_forward_payload", [[PORT_FORWARD_PLEX]])
+@pytest.mark.parametrize(("traffic_rule_payload"), [([TRAFFIC_RULE])])
 @pytest.mark.parametrize("wlan_payload", [[WLAN]])
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -1661,6 +1618,7 @@ async def test_hub_state_change(
         "switch.plug_outlet_1",
         "switch.block_media_streaming",
         "switch.unifi_network_plex",
+        "switch.unifi_network_test_traffic_rule",
         "switch.ssid_1",
     )
     for entity_id in entity_ids: