From 22e62f42e6efe7d50a587709a2ea3e54e35cfc13 Mon Sep 17 00:00:00 2001 From: Jan-Philipp Benecke Date: Sun, 18 Feb 2024 23:47:10 +0100 Subject: [PATCH] Use entity & device registry fixtures in unifi tests (#110889) --- tests/components/unifi/conftest.py | 5 +-- tests/components/unifi/test_button.py | 16 ++++--- tests/components/unifi/test_controller.py | 6 ++- tests/components/unifi/test_device_tracker.py | 10 +++-- tests/components/unifi/test_image.py | 8 ++-- tests/components/unifi/test_sensor.py | 42 +++++++++++-------- tests/components/unifi/test_services.py | 25 ++++++----- tests/components/unifi/test_switch.py | 34 ++++++++------- 8 files changed, 85 insertions(+), 61 deletions(-) diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index 0cddd505cd4..7337ddd99cb 100644 --- a/tests/components/unifi/conftest.py +++ b/tests/components/unifi/conftest.py @@ -108,9 +108,8 @@ def mock_discovery(): @pytest.fixture -def mock_device_registry(hass): +def mock_device_registry(hass, device_registry: dr.DeviceRegistry): """Mock device registry.""" - dev_reg = dr.async_get(hass) config_entry = MockConfigEntry(domain="something_else") config_entry.add_to_hass(hass) @@ -126,7 +125,7 @@ def mock_device_registry(hass): "00:00:00:00:02:02", ) ): - dev_reg.async_get_or_create( + device_registry.async_get_or_create( name=f"Device {idx}", config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, device)}, diff --git a/tests/components/unifi/test_button.py b/tests/components/unifi/test_button.py index 8e6dce71160..11d0e2c3fae 100644 --- a/tests/components/unifi/test_button.py +++ b/tests/components/unifi/test_button.py @@ -12,7 +12,10 @@ from tests.test_util.aiohttp import AiohttpClientMocker async def test_restart_device_button( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + websocket_mock, ) -> None: """Test restarting device button.""" config_entry = await setup_unifi_integration( @@ -37,8 +40,7 @@ async def test_restart_device_button( assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("button.switch_restart") + ent_reg_entry = entity_registry.async_get("button.switch_restart") assert ent_reg_entry.unique_id == "device_restart-00:00:00:00:01:01" assert ent_reg_entry.entity_category is EntityCategory.CONFIG @@ -78,7 +80,10 @@ async def test_restart_device_button( async def test_power_cycle_poe( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + websocket_mock, ) -> None: """Test restarting device button.""" config_entry = await setup_unifi_integration( @@ -119,8 +124,7 @@ async def test_power_cycle_poe( assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 2 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("button.switch_port_1_power_cycle") + ent_reg_entry = entity_registry.async_get("button.switch_port_1_power_cycle") assert ent_reg_entry.unique_id == "power_cycle-00:00:00:00:01:01_1" assert ent_reg_entry.entity_category is EntityCategory.CONFIG diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 54c1d055b70..17ae018d67f 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -242,7 +242,9 @@ async def setup_unifi_integration( async def test_controller_setup( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Successful setup.""" with patch( @@ -278,7 +280,7 @@ async def test_controller_setup( assert controller.signal_options_update == "unifi-options-1" assert controller.signal_heartbeat_missed == "unifi-heartbeat-missed" - device_entry = dr.async_get(hass).async_get_or_create( + device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, identifiers={(UNIFI_DOMAIN, config_entry.unique_id)}, ) diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 34d43129a94..16bd7adf144 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -906,7 +906,10 @@ async def test_option_ignore_wired_bug( async def test_restoring_client( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_device_registry + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + mock_device_registry, ) -> None: """Verify clients are restored from clients_all if they ever was registered to entity registry.""" client = { @@ -939,15 +942,14 @@ async def test_restoring_client( entry_id="1", ) - registry = er.async_get(hass) - registry.async_get_or_create( # Unique ID updated + entity_registry.async_get_or_create( # Unique ID updated TRACKER_DOMAIN, UNIFI_DOMAIN, f'{restored["mac"]}-site_id', suggested_object_id=restored["hostname"], config_entry=config_entry, ) - registry.async_get_or_create( # Unique ID already updated + entity_registry.async_get_or_create( # Unique ID already updated TRACKER_DOMAIN, UNIFI_DOMAIN, f'site_id-{client["mac"]}', diff --git a/tests/components/unifi/test_image.py b/tests/components/unifi/test_image.py index 92879f5ad14..74e1e9b97b3 100644 --- a/tests/components/unifi/test_image.py +++ b/tests/components/unifi/test_image.py @@ -60,6 +60,7 @@ WLAN = { async def test_wlan_qr_code( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, hass_client: ClientSessionGenerator, snapshot: SnapshotAssertion, @@ -70,14 +71,15 @@ async def test_wlan_qr_code( 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") + 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 - ent_reg.async_update_entity(entity_id="image.ssid_1_qr_code", disabled_by=None) + entity_registry.async_update_entity( + entity_id="image.ssid_1_qr_code", disabled_by=None + ) await hass.async_block_till_done() async_fire_time_changed( diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 9ebdd207b54..c89e193d5b7 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -457,6 +457,7 @@ async def test_bandwidth_sensors( ) async def test_uptime_sensors( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, freezer: FrozenDateTimeFactory, mock_unifi_websocket, @@ -492,9 +493,8 @@ async def test_uptime_sensors( assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1 assert hass.states.get("sensor.client1_uptime").state == "2021-01-01T01:00:00+00:00" - ent_reg = er.async_get(hass) assert ( - ent_reg.async_get("sensor.client1_uptime").entity_category + entity_registry.async_get("sensor.client1_uptime").entity_category is EntityCategory.DIAGNOSTIC ) @@ -606,6 +606,7 @@ async def test_remove_sensors( async def test_poe_port_switches( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket, websocket_mock, @@ -614,13 +615,12 @@ async def test_poe_port_switches( await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1]) assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("sensor.mock_name_port_1_poe_power") + ent_reg_entry = entity_registry.async_get("sensor.mock_name_port_1_poe_power") 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_registry.async_update_entity( entity_id="sensor.mock_name_port_1_poe_power", disabled_by=None ) await hass.async_block_till_done() @@ -681,6 +681,7 @@ async def test_poe_port_switches( async def test_wlan_client_sensors( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket, websocket_mock, @@ -716,8 +717,7 @@ async def test_wlan_client_sensors( assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("sensor.ssid_1") + ent_reg_entry = entity_registry.async_get("sensor.ssid_1") assert ent_reg_entry.unique_id == "wlan_clients-012345678910111213141516" assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC @@ -824,6 +824,7 @@ async def test_wlan_client_sensors( ) async def test_outlet_power_readings( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket, entity_id: str, @@ -838,8 +839,7 @@ async def test_outlet_power_readings( assert len(hass.states.async_all()) == 11 assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 5 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get(f"sensor.{entity_id}") + ent_reg_entry = entity_registry.async_get(f"sensor.{entity_id}") assert ent_reg_entry.unique_id == expected_unique_id assert ent_reg_entry.entity_category is EntityCategory.DIAGNOSTIC @@ -859,7 +859,10 @@ async def test_outlet_power_readings( async def test_device_uptime( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + mock_unifi_websocket, ) -> None: """Verify that uptime sensors are working as expected.""" device = { @@ -887,9 +890,8 @@ async def test_device_uptime( assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2 assert hass.states.get("sensor.device_uptime").state == "2021-01-01T01:00:00+00:00" - ent_reg = er.async_get(hass) assert ( - ent_reg.async_get("sensor.device_uptime").entity_category + entity_registry.async_get("sensor.device_uptime").entity_category is EntityCategory.DIAGNOSTIC ) @@ -915,7 +917,10 @@ async def test_device_uptime( async def test_device_temperature( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + mock_unifi_websocket, ) -> None: """Verify that temperature sensors are working as expected.""" device = { @@ -943,9 +948,8 @@ async def test_device_temperature( assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3 assert hass.states.get("sensor.device_temperature").state == "30" - ent_reg = er.async_get(hass) assert ( - ent_reg.async_get("sensor.device_temperature").entity_category + entity_registry.async_get("sensor.device_temperature").entity_category is EntityCategory.DIAGNOSTIC ) @@ -956,7 +960,10 @@ async def test_device_temperature( async def test_device_state( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + mock_unifi_websocket, ) -> None: """Verify that state sensors are working as expected.""" device = { @@ -983,9 +990,8 @@ async def test_device_state( await setup_unifi_integration(hass, aioclient_mock, devices_response=[device]) assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3 - ent_reg = er.async_get(hass) assert ( - ent_reg.async_get("sensor.device_state").entity_category + entity_registry.async_get("sensor.device_state").entity_category is EntityCategory.DIAGNOSTIC ) diff --git a/tests/components/unifi/test_services.py b/tests/components/unifi/test_services.py index 16a7b46c889..4013da11a5a 100644 --- a/tests/components/unifi/test_services.py +++ b/tests/components/unifi/test_services.py @@ -52,7 +52,9 @@ async def test_service_setup_and_unload_not_called_if_multiple_integrations_dete async def test_reconnect_client( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Verify call to reconnect client is performed as expected.""" clients = [ @@ -71,7 +73,6 @@ async def test_reconnect_client( f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", ) - device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])}, @@ -104,14 +105,15 @@ async def test_reconnect_non_existant_device( async def test_reconnect_device_without_mac( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Verify no call is made if device does not have a known mac.""" config_entry = await setup_unifi_integration(hass, aioclient_mock) aioclient_mock.clear_requests() - device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={("other connection", "not mac")}, @@ -127,7 +129,9 @@ async def test_reconnect_device_without_mac( async def test_reconnect_client_controller_unavailable( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Verify no call is made if controller is unavailable.""" clients = [ @@ -147,7 +151,6 @@ async def test_reconnect_client_controller_unavailable( f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", ) - device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])}, @@ -163,14 +166,15 @@ async def test_reconnect_client_controller_unavailable( async def test_reconnect_client_unknown_mac( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Verify no call is made if trying to reconnect a mac unknown to controller.""" config_entry = await setup_unifi_integration(hass, aioclient_mock) aioclient_mock.clear_requests() - device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to controller")}, @@ -186,7 +190,9 @@ async def test_reconnect_client_unknown_mac( async def test_reconnect_wired_client( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Verify no call is made if client is wired.""" clients = [ @@ -201,7 +207,6 @@ async def test_reconnect_wired_client( aioclient_mock.clear_requests() - device_registry = dr.async_get(hass) device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, clients[0]["mac"])}, diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 6a9e58b6f76..bdeab83f6aa 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -816,7 +816,9 @@ async def test_not_admin( async def test_switches( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Test the update_items function with some clients.""" config_entry = await setup_unifi_integration( @@ -852,9 +854,10 @@ async def test_switches( assert dpi_switch.state == "on" assert dpi_switch.attributes["icon"] == "mdi:network" - ent_reg = er.async_get(hass) for entry_id in ("switch.block_client_1", "switch.block_media_streaming"): - assert ent_reg.async_get(entry_id).entity_category is EntityCategory.CONFIG + assert ( + entity_registry.async_get(entry_id).entity_category is EntityCategory.CONFIG + ) # Block and unblock client aioclient_mock.clear_requests() @@ -1326,6 +1329,7 @@ async def test_option_remove_switches( async def test_poe_port_switches( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket, websocket_mock, @@ -1338,16 +1342,15 @@ async def test_poe_port_switches( assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("switch.mock_name_port_1_poe") + 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 - ent_reg.async_update_entity( + entity_registry.async_update_entity( entity_id="switch.mock_name_port_1_poe", disabled_by=None ) - ent_reg.async_update_entity( + entity_registry.async_update_entity( entity_id="switch.mock_name_port_2_poe", disabled_by=None ) await hass.async_block_till_done() @@ -1438,6 +1441,7 @@ async def test_poe_port_switches( async def test_wlan_switches( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket, websocket_mock, @@ -1450,8 +1454,7 @@ async def test_wlan_switches( assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("switch.ssid_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 @@ -1507,6 +1510,7 @@ async def test_wlan_switches( async def test_port_forwarding_switches( hass: HomeAssistant, + entity_registry: er.EntityRegistry, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket, websocket_mock, @@ -1531,8 +1535,7 @@ async def test_port_forwarding_switches( assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 - ent_reg = er.async_get(hass) - ent_reg_entry = ent_reg.async_get("switch.unifi_network_plex") + 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 @@ -1594,7 +1597,9 @@ async def test_port_forwarding_switches( async def test_updating_unique_id( - hass: HomeAssistant, aioclient_mock: AiohttpClientMocker + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, ) -> None: """Verify outlet control and poe control unique ID update works.""" poe_device = { @@ -1637,15 +1642,14 @@ async def test_updating_unique_id( entry_id="1", ) - registry = er.async_get(hass) - registry.async_get_or_create( + entity_registry.async_get_or_create( SWITCH_DOMAIN, UNIFI_DOMAIN, f'{poe_device["mac"]}-poe-1', suggested_object_id="switch_port_1_poe", config_entry=config_entry, ) - registry.async_get_or_create( + entity_registry.async_get_or_create( SWITCH_DOMAIN, UNIFI_DOMAIN, f'{OUTLET_UP1["mac"]}-outlet-1',