Adjust entity registry access in tests (1) (#88950)

tests/components/aladdin_connect/test_sensor.py

@@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 from homeassistant.components.aladdin_connect.const import DOMAIN
 from homeassistant.components.aladdin_connect.cover import SCAN_INTERVAL
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util.dt import utcnow
 
 from tests.common import MockConfigEntry, async_fire_time_changed
@@ -28,6 +28,7 @@ RELOAD_AFTER_UPDATE_DELAY = timedelta(seconds=31)
 async def test_sensors(
     hass: HomeAssistant,
     mock_aladdinconnect_api: MagicMock,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test Sensors for AladdinConnect."""
     config_entry = MockConfigEntry(
@@ -46,12 +47,11 @@ async def test_sensors(
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-        registry = entity_registry.async_get(hass)
+        entry = entity_registry.async_get("sensor.home_battery_level")
-        entry = registry.async_get("sensor.home_battery_level")
         assert entry
         assert entry.disabled
-        assert entry.disabled_by is entity_registry.RegistryEntryDisabler.INTEGRATION
+        assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION
-        update_entry = registry.async_update_entity(
+        update_entry = entity_registry.async_update_entity(
             entry.entity_id, **{"disabled_by": None}
         )
         await hass.async_block_till_done()
@@ -68,12 +68,12 @@ async def test_sensors(
         state = hass.states.get("sensor.home_battery_level")
         assert state
 
-        entry = registry.async_get("sensor.home_wi_fi_rssi")
+        entry = entity_registry.async_get("sensor.home_wi_fi_rssi")
         await hass.async_block_till_done()
         assert entry
         assert entry.disabled
-        assert entry.disabled_by is entity_registry.RegistryEntryDisabler.INTEGRATION
+        assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION
-        update_entry = registry.async_update_entity(
+        update_entry = entity_registry.async_update_entity(
             entry.entity_id, **{"disabled_by": None}
         )
         await hass.async_block_till_done()
@@ -82,7 +82,7 @@ async def test_sensors(
         state = hass.states.get("sensor.home_wi_fi_rssi")
         assert state is None
 
-        update_entry = registry.async_update_entity(
+        update_entry = entity_registry.async_update_entity(
             entry.entity_id, **{"disabled_by": None}
         )
         await hass.async_block_till_done()
@@ -99,6 +99,7 @@ async def test_sensors(
 async def test_sensors_model_01(
     hass: HomeAssistant,
     mock_aladdinconnect_api: MagicMock,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test Sensors for AladdinConnect."""
     config_entry = MockConfigEntry(
@@ -120,20 +121,19 @@ async def test_sensors_model_01(
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-        registry = entity_registry.async_get(hass)
+        entry = entity_registry.async_get("sensor.home_battery_level")
-        entry = registry.async_get("sensor.home_battery_level")
         assert entry
         assert entry.disabled is False
         assert entry.disabled_by is None
         state = hass.states.get("sensor.home_battery_level")
         assert state
 
-        entry = registry.async_get("sensor.home_wi_fi_rssi")
+        entry = entity_registry.async_get("sensor.home_wi_fi_rssi")
         await hass.async_block_till_done()
         assert entry
         assert entry.disabled
-        assert entry.disabled_by is entity_registry.RegistryEntryDisabler.INTEGRATION
+        assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION
-        update_entry = registry.async_update_entity(
+        update_entry = entity_registry.async_update_entity(
             entry.entity_id, **{"disabled_by": None}
         )
         await hass.async_block_till_done()
@@ -142,7 +142,7 @@ async def test_sensors_model_01(
         state = hass.states.get("sensor.home_wi_fi_rssi")
         assert state is None
 
-        update_entry = registry.async_update_entity(
+        update_entry = entity_registry.async_update_entity(
             entry.entity_id, **{"disabled_by": None}
         )
         await hass.async_block_till_done()
@@ -155,7 +155,7 @@ async def test_sensors_model_01(
         state = hass.states.get("sensor.home_wi_fi_rssi")
         assert state
 
-        entry = registry.async_get("sensor.home_ble_strength")
+        entry = entity_registry.async_get("sensor.home_ble_strength")
         await hass.async_block_till_done()
         assert entry
         assert entry.disabled is False

tests/components/broadlink/test_config_flow.py

@@ -10,7 +10,7 @@ from homeassistant import config_entries
 from homeassistant.components import dhcp
 from homeassistant.components.broadlink.const import DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 
 from . import get_device
 
@@ -838,7 +838,7 @@ async def test_dhcp_can_finish(hass: HomeAssistant) -> None:
             data=dhcp.DhcpServiceInfo(
                 hostname="broadlink",
                 ip="1.2.3.4",
-                macaddress=device_registry.format_mac(device.mac),
+                macaddress=dr.format_mac(device.mac),
             ),
         )
         await hass.async_block_till_done()
@@ -932,7 +932,7 @@ async def test_dhcp_device_not_supported(hass: HomeAssistant) -> None:
             data=dhcp.DhcpServiceInfo(
                 hostname="broadlink",
                 ip=device.host,
-                macaddress=device_registry.format_mac(device.mac),
+                macaddress=dr.format_mac(device.mac),
             ),
         )
 

tests/components/coinbase/test_init.py

@@ -9,7 +9,7 @@ from homeassistant.components.coinbase.const import (
     DOMAIN,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 
 from .common import (
     init_mock_coinbase,
@@ -49,7 +49,9 @@ async def test_unload_entry(hass: HomeAssistant) -> None:
     assert not hass.data.get(DOMAIN)
 
 
-async def test_option_updates(hass: HomeAssistant) -> None:
+async def test_option_updates(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test handling option updates."""
 
     with patch(
@@ -75,9 +77,8 @@ async def test_option_updates(hass: HomeAssistant) -> None:
         )
         await hass.async_block_till_done()
 
-        registry = entity_registry.async_get(hass)
+        entities = er.async_entries_for_config_entry(
-        entities = entity_registry.async_entries_for_config_entry(
+            entity_registry, config_entry.entry_id
-            registry, config_entry.entry_id
         )
         assert len(entities) == 4
         currencies = [
@@ -106,9 +107,8 @@ async def test_option_updates(hass: HomeAssistant) -> None:
         )
         await hass.async_block_till_done()
 
-        registry = entity_registry.async_get(hass)
+        entities = er.async_entries_for_config_entry(
-        entities = entity_registry.async_entries_for_config_entry(
+            entity_registry, config_entry.entry_id
-            registry, config_entry.entry_id
         )
         assert len(entities) == 2
         currencies = [
@@ -127,7 +127,9 @@ async def test_option_updates(hass: HomeAssistant) -> None:
         assert rates == [GOOD_EXCHANGE_RATE]
 
 
-async def test_ignore_vaults_wallets(hass: HomeAssistant) -> None:
+async def test_ignore_vaults_wallets(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test vaults are ignored in wallet sensors."""
 
     with patch(
@@ -142,9 +144,8 @@ async def test_ignore_vaults_wallets(hass: HomeAssistant) -> None:
         config_entry = await init_mock_coinbase(hass, currencies=[GOOD_CURRENCY])
         await hass.async_block_till_done()
 
-        registry = entity_registry.async_get(hass)
+        entities = er.async_entries_for_config_entry(
-        entities = entity_registry.async_entries_for_config_entry(
+            entity_registry, config_entry.entry_id
-            registry, config_entry.entry_id
         )
         assert len(entities) == 1
         entity = entities[0]

tests/components/command_line/test_binary_sensor.py

@@ -9,7 +9,7 @@ from homeassistant import setup
 from homeassistant.components.binary_sensor import DOMAIN
 from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 
 
 async def setup_test_entity(hass: HomeAssistant, config_dict: dict[str, Any]) -> None:
@@ -72,7 +72,9 @@ async def test_sensor_off(hass: HomeAssistant) -> None:
     assert entity_state.state == STATE_OFF
 
 
-async def test_unique_id(hass: HomeAssistant) -> None:
+async def test_unique_id(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test unique_id option and if it only creates one binary sensor per id."""
     assert await setup.async_setup_component(
         hass,
@@ -101,18 +103,12 @@ async def test_unique_id(hass: HomeAssistant) -> None:
 
     assert len(hass.states.async_all()) == 2
 
-    ent_reg = entity_registry.async_get(hass)
+    assert len(entity_registry.entities) == 2
-
+    assert entity_registry.async_get_entity_id(
-    assert len(ent_reg.entities) == 2
+        "binary_sensor", "command_line", "unique"
-    assert (
-        ent_reg.async_get_entity_id("binary_sensor", "command_line", "unique")
-        is not None
     )
-    assert (
+    assert entity_registry.async_get_entity_id(
-        ent_reg.async_get_entity_id(
+        "binary_sensor", "command_line", "not-so-unique-anymore"
-            "binary_sensor", "command_line", "not-so-unique-anymore"
-        )
-        is not None
     )
 
 

tests/components/command_line/test_cover.py

@@ -18,7 +18,7 @@ from homeassistant.const import (
     SERVICE_STOP_COVER,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 import homeassistant.util.dt as dt_util
 
 from tests.common import async_fire_time_changed, get_fixture_path
@@ -171,7 +171,9 @@ async def test_move_cover_failure(
     assert "return code 1" in caplog.text
 
 
-async def test_unique_id(hass: HomeAssistant) -> None:
+async def test_unique_id(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test unique_id option and if it only creates one cover per id."""
     await setup_test_entity(
         hass,
@@ -199,11 +201,8 @@ async def test_unique_id(hass: HomeAssistant) -> None:
 
     assert len(hass.states.async_all()) == 2
 
-    ent_reg = entity_registry.async_get(hass)
+    assert len(entity_registry.entities) == 2
-
+    assert entity_registry.async_get_entity_id("cover", "command_line", "unique")
-    assert len(ent_reg.entities) == 2
+    assert entity_registry.async_get_entity_id(
-    assert ent_reg.async_get_entity_id("cover", "command_line", "unique") is not None
+        "cover", "command_line", "not-so-unique-anymore"
-    assert (
-        ent_reg.async_get_entity_id("cover", "command_line", "not-so-unique-anymore")
-        is not None
     )

tests/components/command_line/test_sensor.py

@@ -9,7 +9,7 @@ import pytest
 from homeassistant import setup
 from homeassistant.components.sensor import DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 
 
 async def setup_test_entities(hass: HomeAssistant, config_dict: dict[str, Any]) -> None:
@@ -260,7 +260,9 @@ async def test_update_with_unnecessary_json_attrs(
     assert "key_three" not in entity_state.attributes
 
 
-async def test_unique_id(hass: HomeAssistant) -> None:
+async def test_unique_id(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test unique_id option and if it only creates one sensor per id."""
     assert await setup.async_setup_component(
         hass,
@@ -289,11 +291,8 @@ async def test_unique_id(hass: HomeAssistant) -> None:
 
     assert len(hass.states.async_all()) == 2
 
-    ent_reg = entity_registry.async_get(hass)
+    assert len(entity_registry.entities) == 2
-
+    assert entity_registry.async_get_entity_id("sensor", "command_line", "unique")
-    assert len(ent_reg.entities) == 2
+    assert entity_registry.async_get_entity_id(
-    assert ent_reg.async_get_entity_id("sensor", "command_line", "unique") is not None
+        "sensor", "command_line", "not-so-unique-anymore"
-    assert (
-        ent_reg.async_get_entity_id("sensor", "command_line", "not-so-unique-anymore")
-        is not None
     )

tests/components/command_line/test_switch.py

@@ -20,7 +20,7 @@ from homeassistant.const import (
     STATE_ON,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 import homeassistant.util.dt as dt_util
 
 from tests.common import async_fire_time_changed
@@ -393,7 +393,9 @@ async def test_no_switches(
     assert "No switches" in caplog.text
 
 
-async def test_unique_id(hass: HomeAssistant) -> None:
+async def test_unique_id(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test unique_id option and if it only creates one switch per id."""
     await setup_test_entity(
         hass,
@@ -418,13 +420,10 @@ async def test_unique_id(hass: HomeAssistant) -> None:
 
     assert len(hass.states.async_all()) == 2
 
-    ent_reg = entity_registry.async_get(hass)
+    assert len(entity_registry.entities) == 2
-
+    assert entity_registry.async_get_entity_id("switch", "command_line", "unique")
-    assert len(ent_reg.entities) == 2
+    assert entity_registry.async_get_entity_id(
-    assert ent_reg.async_get_entity_id("switch", "command_line", "unique") is not None
+        "switch", "command_line", "not-so-unique-anymore"
-    assert (
-        ent_reg.async_get_entity_id("switch", "command_line", "not-so-unique-anymore")
-        is not None
     )
 
 

tests/components/devolo_home_control/test_binary_sensor.py

@@ -12,7 +12,7 @@ from homeassistant.const import (
     EntityCategory,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 
 from . import configure_integration
 from .mocks import (
@@ -24,7 +24,9 @@ from .mocks import (
 
 
 @pytest.mark.usefixtures("mock_zeroconf")
-async def test_binary_sensor(hass: HomeAssistant) -> None:
+async def test_binary_sensor(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test setup and state change of a binary sensor device."""
     entry = configure_integration(hass)
     test_gateway = HomeControlMockBinarySensor()
@@ -44,9 +46,8 @@ async def test_binary_sensor(hass: HomeAssistant) -> None:
     state = hass.states.get(f"{DOMAIN}.test_overload")
     assert state is not None
     assert state.attributes[ATTR_FRIENDLY_NAME] == "Test Overload"
-    er = entity_registry.async_get(hass)
     assert (
-        er.async_get(f"{DOMAIN}.test_overload").entity_category
+        entity_registry.async_get(f"{DOMAIN}.test_overload").entity_category
         == EntityCategory.DIAGNOSTIC
     )
 

tests/components/devolo_home_control/test_sensor.py

@@ -15,7 +15,7 @@ from homeassistant.const import (
     EntityCategory,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 
 from . import configure_integration
 from .mocks import HomeControlMock, HomeControlMockConsumption, HomeControlMockSensor
@@ -43,10 +43,11 @@ async def test_temperature_sensor(hass: HomeAssistant) -> None:
     assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.TEMPERATURE
 
 
-async def test_battery_sensor(hass: HomeAssistant) -> None:
+async def test_battery_sensor(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test setup and state change of a battery sensor device."""
     entry = configure_integration(hass)
-    er = entity_registry.async_get(hass)
     test_gateway = HomeControlMockSensor()
     test_gateway.devices["Test"].battery_level = 25
     with patch(
@@ -63,7 +64,7 @@ async def test_battery_sensor(hass: HomeAssistant) -> None:
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == PERCENTAGE
     assert state.attributes[ATTR_DEVICE_CLASS] == SensorDeviceClass.BATTERY
     assert (
-        er.async_get(f"{DOMAIN}.test_battery_level").entity_category
+        entity_registry.async_get(f"{DOMAIN}.test_battery_level").entity_category
         is EntityCategory.DIAGNOSTIC
     )
 

tests/components/devolo_home_network/test_binary_sensor.py

@@ -17,7 +17,7 @@ from homeassistant.const import (
     EntityCategory,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt
 
 from . import configure_integration
@@ -42,15 +42,13 @@ async def test_binary_sensor_setup(hass: HomeAssistant) -> None:
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_update_attached_to_router(
-    hass: HomeAssistant, mock_device: MockDevice
+    hass: HomeAssistant, mock_device: MockDevice, entity_registry: er.EntityRegistry
 ) -> None:
     """Test state change of a attached_to_router binary sensor device."""
     entry = configure_integration(hass)
     device_name = entry.title.replace(" ", "_").lower()
     state_key = f"{DOMAIN}.{device_name}_{CONNECTED_TO_ROUTER}"
 
-    er = entity_registry.async_get(hass)
-
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
@@ -59,7 +57,10 @@ async def test_update_attached_to_router(
     assert state.state == STATE_OFF
     assert state.attributes[ATTR_FRIENDLY_NAME] == f"{entry.title} Connected to router"
 
-    assert er.async_get(state_key).entity_category == EntityCategory.DIAGNOSTIC
+    assert (
+        entity_registry.async_get(state_key).entity_category
+        == EntityCategory.DIAGNOSTIC
+    )
 
     # Emulate device failure
     mock_device.plcnet.async_get_network_overview = AsyncMock(

tests/components/devolo_home_network/test_device_tracker.py

@@ -17,7 +17,7 @@ from homeassistant.const import (
     UnitOfFrequency,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt
 
 from . import configure_integration
@@ -30,20 +30,21 @@ STATION = CONNECTED_STATIONS[0]
 SERIAL = DISCOVERY_INFO.properties["SN"]
 
 
-async def test_device_tracker(hass: HomeAssistant, mock_device: MockDevice) -> None:
+async def test_device_tracker(
+    hass: HomeAssistant, mock_device: MockDevice, entity_registry: er.EntityRegistry
+) -> None:
     """Test device tracker states."""
     state_key = (
         f"{PLATFORM}.{DOMAIN}_{SERIAL}_{STATION.mac_address.lower().replace(':', '_')}"
     )
     entry = configure_integration(hass)
-    er = entity_registry.async_get(hass)
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
     async_fire_time_changed(hass, dt.utcnow() + LONG_UPDATE_INTERVAL)
     await hass.async_block_till_done()
 
     # Enable entity
-    er.async_update_entity(state_key, disabled_by=None)
+    entity_registry.async_update_entity(state_key, disabled_by=None)
     await hass.async_block_till_done()
     async_fire_time_changed(hass, dt.utcnow() + LONG_UPDATE_INTERVAL)
     await hass.async_block_till_done()
@@ -82,14 +83,15 @@ async def test_device_tracker(hass: HomeAssistant, mock_device: MockDevice) -> N
     await hass.config_entries.async_unload(entry.entry_id)
 
 
-async def test_restoring_clients(hass: HomeAssistant, mock_device: MockDevice) -> None:
+async def test_restoring_clients(
+    hass: HomeAssistant, mock_device: MockDevice, entity_registry: er.EntityRegistry
+) -> None:
     """Test restoring existing device_tracker entities."""
     state_key = (
         f"{PLATFORM}.{DOMAIN}_{SERIAL}_{STATION.mac_address.lower().replace(':', '_')}"
     )
     entry = configure_integration(hass)
-    er = entity_registry.async_get(hass)
+    entity_registry.async_get_or_create(
-    er.async_get_or_create(
         PLATFORM,
         DOMAIN,
         f"{SERIAL}_{STATION.mac_address}",

tests/components/devolo_home_network/test_sensor.py

@@ -11,7 +11,7 @@ from homeassistant.components.devolo_home_network.const import (
 from homeassistant.components.sensor import DOMAIN, SensorStateClass
 from homeassistant.const import ATTR_FRIENDLY_NAME, STATE_UNAVAILABLE, EntityCategory
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt
 
 from . import configure_integration
@@ -78,13 +78,12 @@ async def test_update_connected_wifi_clients(
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_update_neighboring_wifi_networks(
-    hass: HomeAssistant, mock_device: MockDevice
+    hass: HomeAssistant, mock_device: MockDevice, entity_registry: er.EntityRegistry
 ) -> None:
     """Test state change of a neighboring_wifi_networks sensor device."""
     entry = configure_integration(hass)
     device_name = entry.title.replace(" ", "_").lower()
     state_key = f"{DOMAIN}.{device_name}_neighboring_wifi_networks"
-    er = entity_registry.async_get(hass)
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
@@ -95,7 +94,10 @@ async def test_update_neighboring_wifi_networks(
         state.attributes[ATTR_FRIENDLY_NAME]
         == f"{entry.title} Neighboring Wifi networks"
     )
-    assert er.async_get(state_key).entity_category is EntityCategory.DIAGNOSTIC
+    assert (
+        entity_registry.async_get(state_key).entity_category
+        is EntityCategory.DIAGNOSTIC
+    )
 
     # Emulate device failure
     mock_device.device.async_get_wifi_neighbor_access_points = AsyncMock(
@@ -122,13 +124,12 @@ async def test_update_neighboring_wifi_networks(
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_update_connected_plc_devices(
-    hass: HomeAssistant, mock_device: MockDevice
+    hass: HomeAssistant, mock_device: MockDevice, entity_registry: er.EntityRegistry
 ) -> None:
     """Test state change of a connected_plc_devices sensor device."""
     entry = configure_integration(hass)
     device_name = entry.title.replace(" ", "_").lower()
     state_key = f"{DOMAIN}.{device_name}_connected_plc_devices"
-    er = entity_registry.async_get(hass)
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
 
@@ -138,7 +139,10 @@ async def test_update_connected_plc_devices(
     assert (
         state.attributes[ATTR_FRIENDLY_NAME] == f"{entry.title} Connected PLC devices"
     )
-    assert er.async_get(state_key).entity_category is EntityCategory.DIAGNOSTIC
+    assert (
+        entity_registry.async_get(state_key).entity_category
+        is EntityCategory.DIAGNOSTIC
+    )
 
     # Emulate device failure
     mock_device.plcnet.async_get_network_overview = AsyncMock(

tests/components/devolo_home_network/test_switch.py

@@ -21,7 +21,7 @@ from homeassistant.const import (
     EntityCategory,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.update_coordinator import REQUEST_REFRESH_DEFAULT_COOLDOWN
 from homeassistant.util import dt
 
@@ -157,7 +157,9 @@ async def test_update_enable_guest_wifi(
     await hass.config_entries.async_unload(entry.entry_id)
 
 
-async def test_update_enable_leds(hass: HomeAssistant, mock_device: MockDevice) -> None:
+async def test_update_enable_leds(
+    hass: HomeAssistant, mock_device: MockDevice, entity_registry: er.EntityRegistry
+) -> None:
     """Test state change of a enable_leds switch device."""
     entry = configure_integration(hass)
     device_name = entry.title.replace(" ", "_").lower()
@@ -170,8 +172,7 @@ async def test_update_enable_leds(hass: HomeAssistant, mock_device: MockDevice)
     assert state is not None
     assert state.state == STATE_OFF
 
-    er = entity_registry.async_get(hass)
+    assert entity_registry.async_get(state_key).entity_category == EntityCategory.CONFIG
-    assert er.async_get(state_key).entity_category == EntityCategory.CONFIG
 
     # Emulate state change
     mock_device.device.async_get_led_setting.return_value = True

tests/components/dlna_dmr/test_init.py

@@ -5,7 +5,7 @@ from unittest.mock import Mock
 from homeassistant.components import media_player
 from homeassistant.components.dlna_dmr.const import DOMAIN as DLNA_DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
@@ -17,6 +17,7 @@ async def test_resource_lifecycle(
     config_entry_mock: MockConfigEntry,
     ssdp_scanner_mock: Mock,
     dmr_device_mock: Mock,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test that resources are acquired/released as the entity is setup/unloaded."""
     # Set up the config entry
@@ -25,8 +26,8 @@ async def test_resource_lifecycle(
     await hass.async_block_till_done()
 
     # Check the entity is created and working
-    entries = entity_registry.async_entries_for_config_entry(
+    entries = er.async_entries_for_config_entry(
-        entity_registry.async_get(hass), config_entry_mock.entry_id
+        entity_registry, config_entry_mock.entry_id
     )
     assert len(entries) == 1
     entity_id = entries[0].entity_id

tests/components/flipr/test_binary_sensor.py

@@ -5,7 +5,7 @@ from unittest.mock import patch
 from homeassistant.components.flipr.const import CONF_FLIPR_ID, DOMAIN
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as entity_reg
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt as dt_util
 
 from tests.common import MockConfigEntry
@@ -23,7 +23,7 @@ MOCK_FLIPR_MEASURE = {
 }
 
 
-async def test_sensors(hass: HomeAssistant) -> None:
+async def test_sensors(hass: HomeAssistant, entity_registry: er.EntityRegistry) -> None:
     """Test the creation and values of the Flipr binary sensors."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -37,8 +37,6 @@ async def test_sensors(hass: HomeAssistant) -> None:
 
     entry.add_to_hass(hass)
 
-    registry = entity_reg.async_get(hass)
-
     with patch(
         "flipr_api.FliprAPIRestClient.get_pool_measure_latest",
         return_value=MOCK_FLIPR_MEASURE,
@@ -47,7 +45,7 @@ async def test_sensors(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     # Check entity unique_id value that is generated in FliprEntity base class.
-    entity = registry.async_get("binary_sensor.flipr_myfliprid_ph_status")
+    entity = entity_registry.async_get("binary_sensor.flipr_myfliprid_ph_status")
     assert entity.unique_id == "myfliprid-ph_status"
 
     state = hass.states.get("binary_sensor.flipr_myfliprid_ph_status")

tests/components/flipr/test_sensor.py

@@ -15,7 +15,7 @@ from homeassistant.const import (
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as entity_reg
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt as dt_util
 
 from tests.common import MockConfigEntry
@@ -34,7 +34,7 @@ MOCK_FLIPR_MEASURE = {
 }
 
 
-async def test_sensors(hass: HomeAssistant) -> None:
+async def test_sensors(hass: HomeAssistant, entity_registry: er.EntityRegistry) -> None:
     """Test the creation and values of the Flipr sensors."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -48,8 +48,6 @@ async def test_sensors(hass: HomeAssistant) -> None:
 
     entry.add_to_hass(hass)
 
-    registry = entity_reg.async_get(hass)
-
     with patch(
         "flipr_api.FliprAPIRestClient.get_pool_measure_latest",
         return_value=MOCK_FLIPR_MEASURE,
@@ -58,7 +56,7 @@ async def test_sensors(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     # Check entity unique_id value that is generated in FliprEntity base class.
-    entity = registry.async_get("sensor.flipr_myfliprid_red_ox")
+    entity = entity_registry.async_get("sensor.flipr_myfliprid_red_ox")
     assert entity.unique_id == "myfliprid-red_ox"
 
     state = hass.states.get("sensor.flipr_myfliprid_ph")
@@ -104,7 +102,9 @@ async def test_sensors(hass: HomeAssistant) -> None:
     assert state.state == "95.0"
 
 
-async def test_error_flipr_api_sensors(hass: HomeAssistant) -> None:
+async def test_error_flipr_api_sensors(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test the Flipr sensors error."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -118,8 +118,6 @@ async def test_error_flipr_api_sensors(hass: HomeAssistant) -> None:
 
     entry.add_to_hass(hass)
 
-    registry = entity_reg.async_get(hass)
-
     with patch(
         "flipr_api.FliprAPIRestClient.get_pool_measure_latest",
         side_effect=FliprError("Error during flipr data retrieval..."),
@@ -128,5 +126,5 @@ async def test_error_flipr_api_sensors(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     # Check entity is not generated because of the FliprError raised.
-    entity = registry.async_get("sensor.flipr_myfliprid_red_ox")
+    entity = entity_registry.async_get("sensor.flipr_myfliprid_red_ox")
     assert entity is None

tests/components/generic/test_config_flow.py

@@ -34,7 +34,7 @@ from homeassistant.const import (
     HTTP_BASIC_AUTHENTICATION,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 
 from tests.common import MockConfigEntry
 from tests.typing import ClientSessionGenerator
@@ -809,11 +809,11 @@ async def test_reload_on_title_change(hass: HomeAssistant) -> None:
     assert hass.states.get("camera.my_title").attributes["friendly_name"] == "New Title"
 
 
-async def test_migrate_existing_ids(hass: HomeAssistant) -> None:
+async def test_migrate_existing_ids(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test that existing ids are migrated for issue #70568."""
 
-    registry = entity_registry.async_get(hass)
-
     test_data = TESTDATA_OPTIONS.copy()
     test_data[CONF_CONTENT_TYPE] = "image/png"
     old_unique_id = "54321"
@@ -825,7 +825,7 @@ async def test_migrate_existing_ids(hass: HomeAssistant) -> None:
     new_unique_id = mock_entry.entry_id
     mock_entry.add_to_hass(hass)
 
-    entity_entry = registry.async_get_or_create(
+    entity_entry = entity_registry.async_get_or_create(
         "camera",
         DOMAIN,
         old_unique_id,
@@ -838,7 +838,7 @@ async def test_migrate_existing_ids(hass: HomeAssistant) -> None:
     await hass.config_entries.async_setup(mock_entry.entry_id)
     await hass.async_block_till_done()
 
-    entity_entry = registry.async_get(entity_id)
+    entity_entry = entity_registry.async_get(entity_id)
     assert entity_entry.unique_id == new_unique_id
 
 

tests/components/harmony/test_switch.py

@@ -16,7 +16,7 @@ from homeassistant.const import (
     STATE_UNAVAILABLE,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.util import utcnow
 
 from .const import ENTITY_PLAY_MUSIC, ENTITY_REMOTE, ENTITY_WATCH_TV, HUB_NAME
@@ -25,7 +25,11 @@ from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_connection_state_changes(
-    harmony_client, mock_hc, hass: HomeAssistant, mock_write_config
+    harmony_client,
+    mock_hc,
+    hass: HomeAssistant,
+    mock_write_config,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Ensure connection changes are reflected in the switch states."""
     entry = MockConfigEntry(
@@ -41,9 +45,8 @@ async def test_connection_state_changes(
     assert not hass.states.get(ENTITY_PLAY_MUSIC)
 
     # enable switch entities
-    ent_reg = entity_registry.async_get(hass)
+    entity_registry.async_update_entity(ENTITY_WATCH_TV, disabled_by=None)
-    ent_reg.async_update_entity(ENTITY_WATCH_TV, disabled_by=None)
+    entity_registry.async_update_entity(ENTITY_PLAY_MUSIC, disabled_by=None)
-    ent_reg.async_update_entity(ENTITY_PLAY_MUSIC, disabled_by=None)
     await hass.config_entries.async_reload(entry.entry_id)
     await hass.async_block_till_done()
 
@@ -80,7 +83,9 @@ async def test_connection_state_changes(
     assert hass.states.is_state(ENTITY_PLAY_MUSIC, STATE_OFF)
 
 
-async def test_switch_toggles(mock_hc, hass: HomeAssistant, mock_write_config) -> None:
+async def test_switch_toggles(
+    mock_hc, hass: HomeAssistant, mock_write_config, entity_registry: er.EntityRegistry
+) -> None:
     """Ensure calls to the switch modify the harmony state."""
     entry = MockConfigEntry(
         domain=DOMAIN, data={CONF_HOST: "192.0.2.0", CONF_NAME: HUB_NAME}
@@ -91,9 +96,8 @@ async def test_switch_toggles(mock_hc, hass: HomeAssistant, mock_write_config) -
     await hass.async_block_till_done()
 
     # enable switch entities
-    ent_reg = entity_registry.async_get(hass)
+    entity_registry.async_update_entity(ENTITY_WATCH_TV, disabled_by=None)
-    ent_reg.async_update_entity(ENTITY_WATCH_TV, disabled_by=None)
+    entity_registry.async_update_entity(ENTITY_PLAY_MUSIC, disabled_by=None)
-    ent_reg.async_update_entity(ENTITY_PLAY_MUSIC, disabled_by=None)
     await hass.config_entries.async_reload(entry.entry_id)
     await hass.async_block_till_done()
 

tests/components/hassio/test_binary_sensor.py

@@ -6,7 +6,7 @@ import pytest
 
 from homeassistant.components.hassio import DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
@@ -157,7 +157,11 @@ def mock_all(aioclient_mock, request):
     ],
 )
 async def test_binary_sensor(
-    hass: HomeAssistant, entity_id, expected, aioclient_mock: AiohttpClientMocker
+    hass: HomeAssistant,
+    entity_id,
+    expected,
+    aioclient_mock: AiohttpClientMocker,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test hassio OS and addons binary sensor."""
     config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
@@ -176,8 +180,7 @@ async def test_binary_sensor(
     assert hass.states.get(entity_id) is None
 
     # Enable the entity.
-    ent_reg = entity_registry.async_get(hass)
+    entity_registry.async_update_entity(entity_id, disabled_by=None)
-    ent_reg.async_update_entity(entity_id, disabled_by=None)
     await hass.config_entries.async_reload(config_entry.entry_id)
     await hass.async_block_till_done()
 

tests/components/hassio/test_sensor.py

@@ -6,7 +6,7 @@ import pytest
 
 from homeassistant.components.hassio import DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
@@ -158,7 +158,11 @@ def mock_all(aioclient_mock, request):
     ],
 )
 async def test_sensor(
-    hass: HomeAssistant, entity_id, expected, aioclient_mock: AiohttpClientMocker
+    hass: HomeAssistant,
+    entity_id,
+    expected,
+    aioclient_mock: AiohttpClientMocker,
+    entity_registry: er.EntityRegistry,
 ) -> None:
     """Test hassio OS and addons sensor."""
     config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
@@ -177,8 +181,7 @@ async def test_sensor(
     assert hass.states.get(entity_id) is None
 
     # Enable the entity.
-    ent_reg = entity_registry.async_get(hass)
+    entity_registry.async_update_entity(entity_id, disabled_by=None)
-    ent_reg.async_update_entity(entity_id, disabled_by=None)
     await hass.config_entries.async_reload(config_entry.entry_id)
     await hass.async_block_till_done()