Use entity & device registry fixtures in unifi tests (#110889)

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)},

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 = entity_registry.async_get("button.switch_restart")
-    ent_reg_entry = ent_reg.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 = entity_registry.async_get("button.switch_port_1_power_cycle")
-    ent_reg_entry = ent_reg.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
 

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)},
     )

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)
+    entity_registry.async_get_or_create(  # Unique ID updated
-    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"]}',

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 = entity_registry.async_get("image.ssid_1_qr_code")
-    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)
+    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(

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 = entity_registry.async_get("sensor.mock_name_port_1_poe_power")
-    ent_reg_entry = ent_reg.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 = entity_registry.async_get("sensor.ssid_1")
-    ent_reg_entry = ent_reg.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 = entity_registry.async_get(f"sensor.{entity_id}")
-    ent_reg_entry = ent_reg.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
     )
 

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"])},

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 = entity_registry.async_get("switch.mock_name_port_1_poe")
-    ent_reg_entry = ent_reg.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 = entity_registry.async_get("switch.ssid_1")
-    ent_reg_entry = ent_reg.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 = entity_registry.async_get("switch.unifi_network_plex")
-    ent_reg_entry = ent_reg.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)
+    entity_registry.async_get_or_create(
-    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',