From 6551eb168d1f7362500d9213793ecad559255565 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Wed, 8 Feb 2023 16:28:44 +0100 Subject: [PATCH] Add common area, device, entity and issue registry test fixtures (#87694) --- pylint/plugins/hass_enforce_type_hints.py | 4 ++ tests/common.py | 36 +++++++++++- tests/components/rainbird/test_init.py | 2 +- tests/components/search/test_init.py | 71 ++++++++++++++--------- tests/conftest.py | 28 +++++++++ 5 files changed, 109 insertions(+), 32 deletions(-) diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py index b30fa54de78..9fd4442151a 100644 --- a/pylint/plugins/hass_enforce_type_hints.py +++ b/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", diff --git a/tests/common.py b/tests/common.py index 717e6fdd2d9..224add31b8e 100644 --- a/tests/common.py +++ b/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: diff --git a/tests/components/rainbird/test_init.py b/tests/components/rainbird/test_init.py index e8e9f76d312..fa399960230 100644 --- a/tests/components/rainbird/test_init.py +++ b/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" ) diff --git a/tests/components/search/test_init.py b/tests/components/search/test_init.py index cc04680d8a4..82bceffcea5 100644 --- a/tests/components/search/test_init.py +++ b/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}, } diff --git a/tests/conftest.py b/tests/conftest.py index 7e732468f6d..0fc2443718d 100644 --- a/tests/conftest.py +++ b/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)