Adjust entity registry access in helper tests (#88965)

tests/helpers/test_intent.py

@@ -7,10 +7,10 @@ from homeassistant.components.switch import SwitchDeviceClass
 from homeassistant.const import ATTR_FRIENDLY_NAME
 from homeassistant.core import Context, HomeAssistant, State
 from homeassistant.helpers import (
-    area_registry,
+    area_registry as ar,
     config_validation as cv,
-    device_registry,
-    entity_registry,
+    device_registry as dr,
+    entity_registry as er,
     intent,
 )
 from homeassistant.setup import async_setup_component
@@ -24,12 +24,15 @@ class MockIntentHandler(intent.IntentHandler):
         self.slot_schema = slot_schema
 
 
-async def test_async_match_states(hass: HomeAssistant) -> None:
+async def test_async_match_states(
+    hass: HomeAssistant,
+    area_registry: ar.AreaRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
     """Test async_match_state helper."""
-    areas = area_registry.async_get(hass)
-    area_kitchen = areas.async_get_or_create("kitchen")
-    areas.async_update(area_kitchen.id, aliases={"food room"})
-    area_bedroom = areas.async_get_or_create("bedroom")
+    area_kitchen = area_registry.async_get_or_create("kitchen")
+    area_registry.async_update(area_kitchen.id, aliases={"food room"})
+    area_bedroom = area_registry.async_get_or_create("bedroom")
 
     state1 = State(
         "light.kitchen", "on", attributes={ATTR_FRIENDLY_NAME: "kitchen light"}
@@ -39,14 +42,15 @@ async def test_async_match_states(hass: HomeAssistant) -> None:
     )
 
     # Put entities into different areas
-    entities = entity_registry.async_get(hass)
-    entities.async_get_or_create("light", "demo", "1234", suggested_object_id="kitchen")
-    entities.async_update_entity(state1.entity_id, area_id=area_kitchen.id)
+    entity_registry.async_get_or_create(
+        "light", "demo", "1234", suggested_object_id="kitchen"
+    )
+    entity_registry.async_update_entity(state1.entity_id, area_id=area_kitchen.id)
 
-    entities.async_get_or_create(
+    entity_registry.async_get_or_create(
         "switch", "demo", "5678", suggested_object_id="bedroom"
     )
-    entities.async_update_entity(
+    entity_registry.async_update_entity(
         state2.entity_id,
         area_id=area_bedroom.id,
         device_class=SwitchDeviceClass.OUTLET,
@@ -102,17 +106,20 @@ async def test_async_match_states(hass: HomeAssistant) -> None:
     ) == [state2]
 
 
-async def test_match_device_area(hass: HomeAssistant) -> None:
+async def test_match_device_area(
+    hass: HomeAssistant,
+    area_registry: ar.AreaRegistry,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
     """Test async_match_state with a device in an area."""
-    areas = area_registry.async_get(hass)
-    area_kitchen = areas.async_get_or_create("kitchen")
-    area_bedroom = areas.async_get_or_create("bedroom")
+    area_kitchen = area_registry.async_get_or_create("kitchen")
+    area_bedroom = area_registry.async_get_or_create("bedroom")
 
-    devices = device_registry.async_get(hass)
-    kitchen_device = devices.async_get_or_create(
+    kitchen_device = device_registry.async_get_or_create(
         config_entry_id="1234", connections=set(), identifiers={("demo", "id-1234")}
     )
-    devices.async_update_device(kitchen_device.id, area_id=area_kitchen.id)
+    device_registry.async_update_device(kitchen_device.id, area_id=area_kitchen.id)
 
     state1 = State(
         "light.kitchen", "on", attributes={ATTR_FRIENDLY_NAME: "kitchen light"}
@@ -123,12 +130,15 @@ async def test_match_device_area(hass: HomeAssistant) -> None:
     state3 = State(
         "light.living_room", "on", attributes={ATTR_FRIENDLY_NAME: "living room light"}
     )
-    entities = entity_registry.async_get(hass)
-    entities.async_get_or_create("light", "demo", "1234", suggested_object_id="kitchen")
-    entities.async_update_entity(state1.entity_id, device_id=kitchen_device.id)
+    entity_registry.async_get_or_create(
+        "light", "demo", "1234", suggested_object_id="kitchen"
+    )
+    entity_registry.async_update_entity(state1.entity_id, device_id=kitchen_device.id)
 
-    entities.async_get_or_create("light", "demo", "5678", suggested_object_id="bedroom")
-    entities.async_update_entity(state2.entity_id, area_id=area_bedroom.id)
+    entity_registry.async_get_or_create(
+        "light", "demo", "5678", suggested_object_id="bedroom"
+    )
+    entity_registry.async_update_entity(state2.entity_id, area_id=area_bedroom.id)
 
     # Match on area/domain
     assert list(

tests/helpers/test_service.py

@@ -21,8 +21,8 @@ from homeassistant.const import (
 )
 from homeassistant.core import Context, HomeAssistant, ServiceCall
 from homeassistant.helpers import (
-    device_registry as dev_reg,
-    entity_registry as ent_reg,
+    device_registry as dr,
+    entity_registry as er,
     service,
     template,
 )
@@ -96,10 +96,10 @@ def area_mock(hass):
     hass.states.async_set("light.Ceiling", STATE_OFF)
     hass.states.async_set("light.Kitchen", STATE_OFF)
 
-    device_in_area = dev_reg.DeviceEntry(area_id="test-area")
-    device_no_area = dev_reg.DeviceEntry(id="device-no-area-id")
-    device_diff_area = dev_reg.DeviceEntry(area_id="diff-area")
-    device_area_a = dev_reg.DeviceEntry(id="device-area-a-id", area_id="area-a")
+    device_in_area = dr.DeviceEntry(area_id="test-area")
+    device_no_area = dr.DeviceEntry(id="device-no-area-id")
+    device_diff_area = dr.DeviceEntry(area_id="diff-area")
+    device_area_a = dr.DeviceEntry(id="device-area-a-id", area_id="area-a")
 
     mock_device_registry(
         hass,
@@ -111,94 +111,94 @@ def area_mock(hass):
         },
     )
 
-    entity_in_own_area = ent_reg.RegistryEntry(
+    entity_in_own_area = er.RegistryEntry(
         entity_id="light.in_own_area",
         unique_id="in-own-area-id",
         platform="test",
         area_id="own-area",
     )
-    config_entity_in_own_area = ent_reg.RegistryEntry(
+    config_entity_in_own_area = er.RegistryEntry(
         entity_id="light.config_in_own_area",
         unique_id="config-in-own-area-id",
         platform="test",
         area_id="own-area",
         entity_category=EntityCategory.CONFIG,
     )
-    hidden_entity_in_own_area = ent_reg.RegistryEntry(
+    hidden_entity_in_own_area = er.RegistryEntry(
         entity_id="light.hidden_in_own_area",
         unique_id="hidden-in-own-area-id",
         platform="test",
         area_id="own-area",
-        hidden_by=ent_reg.RegistryEntryHider.USER,
+        hidden_by=er.RegistryEntryHider.USER,
     )
-    entity_in_area = ent_reg.RegistryEntry(
+    entity_in_area = er.RegistryEntry(
         entity_id="light.in_area",
         unique_id="in-area-id",
         platform="test",
         device_id=device_in_area.id,
     )
-    config_entity_in_area = ent_reg.RegistryEntry(
+    config_entity_in_area = er.RegistryEntry(
         entity_id="light.config_in_area",
         unique_id="config-in-area-id",
         platform="test",
         device_id=device_in_area.id,
         entity_category=EntityCategory.CONFIG,
     )
-    hidden_entity_in_area = ent_reg.RegistryEntry(
+    hidden_entity_in_area = er.RegistryEntry(
         entity_id="light.hidden_in_area",
         unique_id="hidden-in-area-id",
         platform="test",
         device_id=device_in_area.id,
-        hidden_by=ent_reg.RegistryEntryHider.USER,
+        hidden_by=er.RegistryEntryHider.USER,
     )
-    entity_in_other_area = ent_reg.RegistryEntry(
+    entity_in_other_area = er.RegistryEntry(
         entity_id="light.in_other_area",
         unique_id="in-area-a-id",
         platform="test",
         device_id=device_in_area.id,
         area_id="other-area",
     )
-    entity_assigned_to_area = ent_reg.RegistryEntry(
+    entity_assigned_to_area = er.RegistryEntry(
         entity_id="light.assigned_to_area",
         unique_id="assigned-area-id",
         platform="test",
         device_id=device_in_area.id,
         area_id="test-area",
     )
-    entity_no_area = ent_reg.RegistryEntry(
+    entity_no_area = er.RegistryEntry(
         entity_id="light.no_area",
         unique_id="no-area-id",
         platform="test",
         device_id=device_no_area.id,
     )
-    config_entity_no_area = ent_reg.RegistryEntry(
+    config_entity_no_area = er.RegistryEntry(
         entity_id="light.config_no_area",
         unique_id="config-no-area-id",
         platform="test",
         device_id=device_no_area.id,
         entity_category=EntityCategory.CONFIG,
     )
-    hidden_entity_no_area = ent_reg.RegistryEntry(
+    hidden_entity_no_area = er.RegistryEntry(
         entity_id="light.hidden_no_area",
         unique_id="hidden-no-area-id",
         platform="test",
         device_id=device_no_area.id,
-        hidden_by=ent_reg.RegistryEntryHider.USER,
+        hidden_by=er.RegistryEntryHider.USER,
     )
-    entity_diff_area = ent_reg.RegistryEntry(
+    entity_diff_area = er.RegistryEntry(
         entity_id="light.diff_area",
         unique_id="diff-area-id",
         platform="test",
         device_id=device_diff_area.id,
     )
-    entity_in_area_a = ent_reg.RegistryEntry(
+    entity_in_area_a = er.RegistryEntry(
         entity_id="light.in_area_a",
         unique_id="in-area-a-id",
         platform="test",
         device_id=device_area_a.id,
         area_id="area-a",
     )
-    entity_in_area_b = ent_reg.RegistryEntry(
+    entity_in_area_b = er.RegistryEntry(
         entity_id="light.in_area_b",
         unique_id="in-area-b-id",
         platform="test",
@@ -403,11 +403,12 @@ class TestServiceHelpers(unittest.TestCase):
         assert mock_log.call_count == 3
 
 
-async def test_service_call_entry_id(hass: HomeAssistant) -> None:
+async def test_service_call_entry_id(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Test service call with entity specified by entity registry ID."""
-    registry = ent_reg.async_get(hass)
     calls = async_mock_service(hass, "test_domain", "test_service")
-    entry = registry.async_get_or_create(
+    entry = entity_registry.async_get_or_create(
         "hello", "hue", "1234", suggested_object_id="world"
     )
 
@@ -946,7 +947,7 @@ async def test_domain_control_unauthorized(
     mock_registry(
         hass,
         {
-            "light.kitchen": ent_reg.RegistryEntry(
+            "light.kitchen": er.RegistryEntry(
                 entity_id="light.kitchen",
                 unique_id="kitchen",
                 platform="test_domain",
@@ -987,7 +988,7 @@ async def test_domain_control_admin(
     mock_registry(
         hass,
         {
-            "light.kitchen": ent_reg.RegistryEntry(
+            "light.kitchen": er.RegistryEntry(
                 entity_id="light.kitchen",
                 unique_id="kitchen",
                 platform="test_domain",
@@ -1025,7 +1026,7 @@ async def test_domain_control_no_user(hass: HomeAssistant) -> None:
     mock_registry(
         hass,
         {
-            "light.kitchen": ent_reg.RegistryEntry(
+            "light.kitchen": er.RegistryEntry(
                 entity_id="light.kitchen",
                 unique_id="kitchen",
                 platform="test_domain",
@@ -1213,9 +1214,7 @@ async def test_async_extract_entities_warn_referenced(
 async def test_async_extract_config_entry_ids(hass: HomeAssistant) -> None:
     """Test we can find devices that have no entities."""
 
-    device_no_entities = dev_reg.DeviceEntry(
-        id="device-no-entities", config_entries={"abc"}
-    )
+    device_no_entities = dr.DeviceEntry(id="device-no-entities", config_entries={"abc"})
 
     call = ServiceCall(
         "homeassistant",