Add common area, device, entity and issue registry test fixtures (#87694)

pylint/plugins/hass_enforce_type_hints.py

@@ -93,10 +93,14 @@ _METHOD_MATCH: list[TypeHintMatch] = [
 _TEST_FIXTURES: dict[str, list[str] | str] = {
     "aioclient_mock": "AiohttpClientMocker",
     "aiohttp_client": "ClientSessionGenerator",
+    "area_registry": "AreaRegistry",
     "caplog": "pytest.LogCaptureFixture",
+    "device_registry": "DeviceRegistry",
+    "entity_registry": "EntityRegistry",
     "hass_client": "ClientSessionGenerator",
     "hass_client_no_auth": "ClientSessionGenerator",
     "hass_ws_client": "WebSocketGenerator",
+    "issue_registry": "IssueRegistry",
     "mqtt_client_mock": "MqttMockPahoClient",
     "mqtt_mock": "MqttMockHAClient",
     "mqtt_mock_entry_no_yaml_config": "MqttMockHAClientGenerator",

tests/common.py

@@ -534,7 +534,17 @@ def mock_registry(
     hass: HomeAssistant,
     mock_entries: dict[str, entity_registry.RegistryEntry] | None = None,
 ) -> entity_registry.EntityRegistry:
-    """Mock the Entity Registry."""
+    """Mock the Entity Registry.
+
+    This should only be used if you need to mock/re-stage a clean mocked
+    entity registry in your current hass object. It can be useful to,
+    for example, pre-load the registry with items.
+
+    This mock will thus replace the existing registry in the running hass.
+
+    If you just need to access the existing registry, use the `entity_registry`
+    fixture instead.
+    """
     registry = entity_registry.EntityRegistry(hass)
     if mock_entries is None:
         mock_entries = {}
@@ -549,7 +559,17 @@ def mock_registry(
 def mock_area_registry(
     hass: HomeAssistant, mock_entries: dict[str, area_registry.AreaEntry] | None = None
 ) -> area_registry.AreaRegistry:
-    """Mock the Area Registry."""
+    """Mock the Area Registry.
+
+    This should only be used if you need to mock/re-stage a clean mocked
+    area registry in your current hass object. It can be useful to,
+    for example, pre-load the registry with items.
+
+    This mock will thus replace the existing registry in the running hass.
+
+    If you just need to access the existing registry, use the `area_registry`
+    fixture instead.
+    """
     registry = area_registry.AreaRegistry(hass)
     registry.areas = mock_entries or OrderedDict()
 
@@ -561,7 +581,17 @@ def mock_device_registry(
     hass: HomeAssistant,
     mock_entries: dict[str, device_registry.DeviceEntry] | None = None,
 ) -> device_registry.DeviceRegistry:
-    """Mock the Device Registry."""
+    """Mock the Device Registry.
+
+    This should only be used if you need to mock/re-stage a clean mocked
+    device registry in your current hass object. It can be useful to,
+    for example, pre-load the registry with items.
+
+    This mock will thus replace the existing registry in the running hass.
+
+    If you just need to access the existing registry, use the `device_registry`
+    fixture instead.
+    """
     registry = device_registry.DeviceRegistry(hass)
     registry.devices = device_registry.DeviceRegistryItems()
     if mock_entries is None:

tests/components/rainbird/test_init.py

@@ -109,6 +109,7 @@ async def test_rain_delay_service(
     aioclient_mock: AiohttpClientMocker,
     responses: list[str],
     config_entry: ConfigEntry,
+    issue_registry: ir.IssueRegistry,
 ) -> None:
     """Test calling the rain delay service."""
 
@@ -131,7 +132,6 @@ async def test_rain_delay_service(
 
     assert len(aioclient_mock.mock_calls) == 1
 
-    issue_registry: ir.IssueRegistry = ir.async_get(hass)
     issue = issue_registry.async_get_issue(
         domain=DOMAIN, issue_id="deprecated_raindelay"
     )

tests/components/search/test_init.py

@@ -1,5 +1,6 @@
 """Tests for Search integration."""
 from homeassistant.components import search
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers import (
     area_registry as ar,
     device_registry as dr,
@@ -24,27 +25,28 @@ MOCK_ENTITY_SOURCES = {
 }
 
 
-async def test_search(hass):
+async def test_search(
+    hass: HomeAssistant,
+    area_registry: ar.AreaRegistry,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
     """Test that search works."""
-    area_reg = ar.async_get(hass)
-    device_reg = dr.async_get(hass)
-    entity_reg = er.async_get(hass)
-
-    living_room_area = area_reg.async_create("Living Room")
+    living_room_area = area_registry.async_create("Living Room")
 
     # Light strip with 2 lights.
     wled_config_entry = MockConfigEntry(domain="wled")
     wled_config_entry.add_to_hass(hass)
 
-    wled_device = device_reg.async_get_or_create(
+    wled_device = device_registry.async_get_or_create(
         config_entry_id=wled_config_entry.entry_id,
         name="Light Strip",
         identifiers=({"wled", "wled-1"}),
     )
 
-    device_reg.async_update_device(wled_device.id, area_id=living_room_area.id)
+    device_registry.async_update_device(wled_device.id, area_id=living_room_area.id)
 
-    wled_segment_1_entity = entity_reg.async_get_or_create(
+    wled_segment_1_entity = entity_registry.async_get_or_create(
         "light",
         "wled",
         "wled-1-seg-1",
@@ -52,7 +54,7 @@ async def test_search(hass):
         config_entry=wled_config_entry,
         device_id=wled_device.id,
     )
-    wled_segment_2_entity = entity_reg.async_get_or_create(
+    wled_segment_2_entity = entity_registry.async_get_or_create(
         "light",
         "wled",
         "wled-1-seg-2",
@@ -74,20 +76,20 @@ async def test_search(hass):
     }
 
     # Non related info.
-    kitchen_area = area_reg.async_create("Kitchen")
+    kitchen_area = area_registry.async_create("Kitchen")
 
     hue_config_entry = MockConfigEntry(domain="hue")
     hue_config_entry.add_to_hass(hass)
 
-    hue_device = device_reg.async_get_or_create(
+    hue_device = device_registry.async_get_or_create(
         config_entry_id=hue_config_entry.entry_id,
         name="Light Strip",
         identifiers=({"hue", "hue-1"}),
     )
 
-    device_reg.async_update_device(hue_device.id, area_id=kitchen_area.id)
+    device_registry.async_update_device(hue_device.id, area_id=kitchen_area.id)
 
-    hue_segment_1_entity = entity_reg.async_get_or_create(
+    hue_segment_1_entity = entity_registry.async_get_or_create(
         "light",
         "hue",
         "hue-1-seg-1",
@@ -95,7 +97,7 @@ async def test_search(hass):
         config_entry=hue_config_entry,
         device_id=hue_device.id,
     )
-    hue_segment_2_entity = entity_reg.async_get_or_create(
+    hue_segment_2_entity = entity_registry.async_get_or_create(
         "light",
         "hue",
         "hue-1-seg-2",
@@ -262,7 +264,9 @@ async def test_search(hass):
         ("automation", "automation.wled_entity"),
         ("automation", "automation.wled_device"),
     ):
-        searcher = search.Searcher(hass, device_reg, entity_reg, entity_sources)
+        searcher = search.Searcher(
+            hass, device_registry, entity_registry, entity_sources
+        )
         results = searcher.async_search(search_type, search_id)
         # Add the item we searched for, it's omitted from results
         results.setdefault(search_type, set()).add(search_id)
@@ -295,7 +299,9 @@ async def test_search(hass):
         ("scene", "scene.scene_wled_hue"),
         ("group", "group.wled_hue"),
     ):
-        searcher = search.Searcher(hass, device_reg, entity_reg, entity_sources)
+        searcher = search.Searcher(
+            hass, device_registry, entity_registry, entity_sources
+        )
         results = searcher.async_search(search_type, search_id)
         # Add the item we searched for, it's omitted from results
         results.setdefault(search_type, set()).add(search_id)
@@ -317,7 +323,9 @@ async def test_search(hass):
         ("script", "script.non_existing"),
         ("automation", "automation.non_existing"),
     ):
-        searcher = search.Searcher(hass, device_reg, entity_reg, entity_sources)
+        searcher = search.Searcher(
+            hass, device_registry, entity_registry, entity_sources
+        )
         assert searcher.async_search(search_type, search_id) == {}
 
     # Test search of templated script. We can't find referenced areas, devices or
@@ -329,22 +337,25 @@ async def test_search(hass):
         ("device", "script.script_with_templated_services"),
         ("entity", "script.script_with_templated_services"),
     ):
-        searcher = search.Searcher(hass, device_reg, entity_reg, entity_sources)
+        searcher = search.Searcher(
+            hass, device_registry, entity_registry, entity_sources
+        )
         assert searcher.async_search(search_type, search_id) == {}
 
-    searcher = search.Searcher(hass, device_reg, entity_reg, entity_sources)
+    searcher = search.Searcher(hass, device_registry, entity_registry, entity_sources)
     assert searcher.async_search("entity", "light.wled_config_entry_source") == {
         "config_entry": {wled_config_entry.entry_id},
     }
 
 
-async def test_area_lookup(hass):
+async def test_area_lookup(
+    hass: HomeAssistant,
+    area_registry: ar.AreaRegistry,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
     """Test area based lookup."""
-    area_reg = ar.async_get(hass)
-    device_reg = dr.async_get(hass)
-    entity_reg = er.async_get(hass)
-
-    living_room_area = area_reg.async_create("Living Room")
+    living_room_area = area_registry.async_create("Living Room")
 
     await async_setup_component(
         hass,
@@ -384,13 +395,17 @@ async def test_area_lookup(hass):
         },
     )
 
-    searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES)
+    searcher = search.Searcher(
+        hass, device_registry, entity_registry, MOCK_ENTITY_SOURCES
+    )
     assert searcher.async_search("area", living_room_area.id) == {
         "script": {"script.wled"},
         "automation": {"automation.area_turn_on"},
     }
 
-    searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES)
+    searcher = search.Searcher(
+        hass, device_registry, entity_registry, MOCK_ENTITY_SOURCES
+    )
     assert searcher.async_search("automation", "automation.area_turn_on") == {
         "area": {living_room_area.id},
     }

tests/conftest.py

@@ -45,8 +45,12 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import HASSIO_USER_NAME
 from homeassistant.core import CoreState, HomeAssistant
 from homeassistant.helpers import (
+    area_registry as ar,
     config_entry_oauth2_flow,
+    device_registry as dr,
+    entity_registry as er,
     event,
+    issue_registry as ir,
     recorder as recorder_helper,
 )
 from homeassistant.helpers.json import json_loads
@@ -1254,3 +1258,27 @@ def mock_bleak_scanner_start():
 @pytest.fixture(name="mock_bluetooth")
 def mock_bluetooth(mock_bleak_scanner_start, mock_bluetooth_adapters):
     """Mock out bluetooth from starting."""
+
+
+@pytest.fixture
+def area_registry(hass: HomeAssistant) -> ar.AreaRegistry:
+    """Return the area registry from the current hass instance."""
+    return ar.async_get(hass)
+
+
+@pytest.fixture
+def device_registry(hass: HomeAssistant) -> dr.DeviceRegistry:
+    """Return the device registry from the current hass instance."""
+    return dr.async_get(hass)
+
+
+@pytest.fixture
+def entity_registry(hass: HomeAssistant) -> er.EntityRegistry:
+    """Return the entity registry from the current hass instance."""
+    return er.async_get(hass)
+
+
+@pytest.fixture
+def issue_registry(hass: HomeAssistant) -> ir.IssueRegistry:
+    """Return the issue registry from the current hass instance."""
+    return ir.async_get(hass)