diff --git a/tests/components/tasmota/test_common.py b/tests/components/tasmota/test_common.py index 973ecd3c890..74e8d2a5e59 100644 --- a/tests/components/tasmota/test_common.py +++ b/tests/components/tasmota/test_common.py @@ -20,6 +20,7 @@ from hatasmota.utils import ( from homeassistant.components.tasmota.const import DEFAULT_PREFIX from homeassistant.const import STATE_UNAVAILABLE +from homeassistant.helpers import device_registry as dr, entity_registry as er from tests.common import async_fire_mqtt_message @@ -363,8 +364,8 @@ async def help_test_discovery_removal( name="Test", ): """Test removal of discovered entity.""" - device_reg = await hass.helpers.device_registry.async_get_registry() - entity_reg = await hass.helpers.entity_registry.async_get_registry() + device_reg = dr.async_get(hass) + entity_reg = er.async_get(hass) data1 = json.dumps(config1) data2 = json.dumps(config2) @@ -470,8 +471,8 @@ async def help_test_discovery_device_remove( hass, mqtt_mock, domain, unique_id, config, sensor_config=None ): """Test domain entity is removed when device is removed.""" - device_reg = await hass.helpers.device_registry.async_get_registry() - entity_reg = await hass.helpers.entity_registry.async_get_registry() + device_reg = dr.async_get(hass) + entity_reg = er.async_get(hass) config = copy.deepcopy(config) @@ -502,7 +503,7 @@ async def help_test_entity_id_update_subscriptions( hass, mqtt_mock, domain, config, topics=None, sensor_config=None, entity_id="test" ): """Test MQTT subscriptions are managed when entity_id is updated.""" - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) config = copy.deepcopy(config) data = json.dumps(config) @@ -548,7 +549,7 @@ async def help_test_entity_id_update_discovery_update( hass, mqtt_mock, domain, config, sensor_config=None, entity_id="test" ): """Test MQTT discovery update after entity_id is updated.""" - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) config = copy.deepcopy(config) data = json.dumps(config) diff --git a/tests/components/tasmota/test_sensor.py b/tests/components/tasmota/test_sensor.py index 6e5160273d6..c08b1b531de 100644 --- a/tests/components/tasmota/test_sensor.py +++ b/tests/components/tasmota/test_sensor.py @@ -17,6 +17,7 @@ from homeassistant import config_entries from homeassistant.components import sensor from homeassistant.components.tasmota.const import DEFAULT_PREFIX from homeassistant.const import ATTR_ASSUMED_STATE, STATE_UNKNOWN +from homeassistant.helpers import entity_registry as er from homeassistant.util import dt from .test_common import ( @@ -226,7 +227,7 @@ async def test_indexed_sensor_state_via_mqtt(hass, mqtt_mock, setup_tasmota): @pytest.mark.parametrize("status_sensor_disabled", [False]) async def test_status_sensor_state_via_mqtt(hass, mqtt_mock, setup_tasmota): """Test state update via MQTT.""" - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) # Pre-enable the status sensor entity_reg.async_get_or_create( @@ -285,7 +286,7 @@ async def test_status_sensor_state_via_mqtt(hass, mqtt_mock, setup_tasmota): @pytest.mark.parametrize("status_sensor_disabled", [False]) async def test_single_shot_status_sensor_state_via_mqtt(hass, mqtt_mock, setup_tasmota): """Test state update via MQTT.""" - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) # Pre-enable the status sensor entity_reg.async_get_or_create( @@ -369,7 +370,7 @@ async def test_restart_time_status_sensor_state_via_mqtt( hass, mqtt_mock, setup_tasmota ): """Test state update via MQTT.""" - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) # Pre-enable the status sensor entity_reg.async_get_or_create( @@ -524,7 +525,7 @@ async def test_indexed_sensor_attributes(hass, mqtt_mock, setup_tasmota): @pytest.mark.parametrize("status_sensor_disabled", [False]) async def test_enable_status_sensor(hass, mqtt_mock, setup_tasmota): """Test enabling status sensor.""" - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) config = copy.deepcopy(DEFAULT_CONFIG) mac = config["mac"] diff --git a/tests/components/timer/test_init.py b/tests/components/timer/test_init.py index 74c3eceeea2..21f48ed6147 100644 --- a/tests/components/timer/test_init.py +++ b/tests/components/timer/test_init.py @@ -39,7 +39,7 @@ from homeassistant.const import ( ) from homeassistant.core import Context, CoreState from homeassistant.exceptions import Unauthorized -from homeassistant.helpers import config_validation as cv, entity_registry +from homeassistant.helpers import config_validation as cv, entity_registry as er from homeassistant.setup import async_setup_component from homeassistant.util.dt import utcnow @@ -248,7 +248,7 @@ async def test_no_initial_state_and_no_restore_state(hass): async def test_config_reload(hass, hass_admin_user, hass_read_only_user): """Test reload service.""" count_start = len(hass.states.async_entity_ids()) - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) _LOGGER.debug("ENTITIES @ start: %s", hass.states.async_entity_ids()) @@ -498,7 +498,7 @@ async def test_ws_delete(hass, hass_ws_client, storage_setup): timer_id = "from_storage" timer_entity_id = f"{DOMAIN}.{DOMAIN}_{timer_id}" - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) state = hass.states.get(timer_entity_id) assert state is not None @@ -525,7 +525,7 @@ async def test_update(hass, hass_ws_client, storage_setup): timer_id = "from_storage" timer_entity_id = f"{DOMAIN}.{DOMAIN}_{timer_id}" - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) state = hass.states.get(timer_entity_id) assert state.attributes[ATTR_FRIENDLY_NAME] == "timer from storage" @@ -554,7 +554,7 @@ async def test_ws_create(hass, hass_ws_client, storage_setup): timer_id = "new_timer" timer_entity_id = f"{DOMAIN}.{timer_id}" - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) state = hass.states.get(timer_entity_id) assert state is None diff --git a/tests/components/traccar/test_init.py b/tests/components/traccar/test_init.py index c7e031b2ca6..0e741751b8a 100644 --- a/tests/components/traccar/test_init.py +++ b/tests/components/traccar/test_init.py @@ -14,6 +14,7 @@ from homeassistant.const import ( STATE_HOME, STATE_NOT_HOME, ) +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.dispatcher import DATA_DISPATCHER from homeassistant.setup import async_setup_component @@ -136,10 +137,10 @@ async def test_enter_and_exit(hass, client, webhook_id): ).state assert STATE_NOT_HOME == state_name - dev_reg = await hass.helpers.device_registry.async_get_registry() + dev_reg = dr.async_get(hass) assert len(dev_reg.devices) == 1 - ent_reg = await hass.helpers.entity_registry.async_get_registry() + ent_reg = er.async_get(hass) assert len(ent_reg.entities) == 1 diff --git a/tests/components/tradfri/test_init.py b/tests/components/tradfri/test_init.py index 0983b5aa22f..da9ae9da146 100644 --- a/tests/components/tradfri/test_init.py +++ b/tests/components/tradfri/test_init.py @@ -2,10 +2,7 @@ from unittest.mock import patch from homeassistant.components import tradfri -from homeassistant.helpers.device_registry import ( - async_entries_for_config_entry, - async_get_registry as async_get_device_registry, -) +from homeassistant.helpers import device_registry as dr from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry @@ -99,8 +96,8 @@ async def test_entry_setup_unload(hass, api_factory, gateway_id): await hass.async_block_till_done() assert setup.call_count == len(tradfri.PLATFORMS) - dev_reg = await async_get_device_registry(hass) - dev_entries = async_entries_for_config_entry(dev_reg, entry.entry_id) + dev_reg = dr.async_get(hass) + dev_entries = dr.async_entries_for_config_entry(dev_reg, entry.entry_id) assert dev_entries dev_entry = dev_entries[0] diff --git a/tests/components/twinkly/test_twinkly.py b/tests/components/twinkly/test_twinkly.py index c8158354195..d4afe02c11b 100644 --- a/tests/components/twinkly/test_twinkly.py +++ b/tests/components/twinkly/test_twinkly.py @@ -1,5 +1,4 @@ """Tests for the integration of a twinly device.""" - from typing import Tuple from unittest.mock import patch @@ -12,6 +11,7 @@ from homeassistant.components.twinkly.const import ( ) from homeassistant.components.twinkly.light import TwinklyLight from homeassistant.core import HomeAssistant +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.device_registry import DeviceEntry from homeassistant.helpers.entity_registry import RegistryEntry @@ -211,8 +211,8 @@ async def _create_entries( assert await hass.config_entries.async_setup(client.id) await hass.async_block_till_done() - device_registry = await hass.helpers.device_registry.async_get_registry() - entity_registry = await hass.helpers.entity_registry.async_get_registry() + device_registry = dr.async_get(hass) + entity_registry = er.async_get(hass) entity_id = entity_registry.async_get_entity_id("light", TWINKLY_DOMAIN, client.id) entity = entity_registry.async_get(entity_id) diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 33dda33be24..2018ae39b64 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -23,7 +23,7 @@ from homeassistant.components.unifi.const import ( DOMAIN as UNIFI_DOMAIN, ) from homeassistant.const import STATE_HOME, STATE_NOT_HOME, STATE_UNAVAILABLE -from homeassistant.helpers import entity_registry +from homeassistant.helpers import device_registry as dr, entity_registry as er import homeassistant.util.dt as dt_util from .test_controller import ENTRY_CONFIG, setup_unifi_integration @@ -335,9 +335,9 @@ async def test_tracked_devices(hass, aioclient_mock, mock_unifi_websocket): await hass.async_block_till_done() # Verify device registry has been updated - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) entry = entity_registry.async_get("device_tracker.device_2") - device_registry = await hass.helpers.device_registry.async_get_registry() + device_registry = dr.async_get(hass) device = device_registry.async_get(entry.device_id) assert device.sw_version == event["version_to"] @@ -949,7 +949,7 @@ async def test_restoring_client(hass, aioclient_mock): entry_id=1, ) - registry = await entity_registry.async_get_registry(hass) + registry = er.async_get(hass) registry.async_get_or_create( TRACKER_DOMAIN, UNIFI_DOMAIN, diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 44120a18ee3..b8f6e2da553 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -16,7 +16,7 @@ from homeassistant.components.unifi.const import ( DOMAIN as UNIFI_DOMAIN, ) from homeassistant.components.unifi.switch import POE_SWITCH -from homeassistant.helpers import entity_registry +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.dispatcher import async_dispatcher_send from .test_controller import ( @@ -799,7 +799,7 @@ async def test_restore_client_succeed(hass, aioclient_mock): entry_id=1, ) - registry = await entity_registry.async_get_registry(hass) + registry = er.async_get(hass) registry.async_get_or_create( SWITCH_DOMAIN, UNIFI_DOMAIN, @@ -891,7 +891,7 @@ async def test_restore_client_no_old_state(hass, aioclient_mock): entry_id=1, ) - registry = await entity_registry.async_get_registry(hass) + registry = er.async_get(hass) registry.async_get_or_create( SWITCH_DOMAIN, UNIFI_DOMAIN, diff --git a/tests/components/vera/test_init.py b/tests/components/vera/test_init.py index 33b6843d7e5..85735d0320e 100644 --- a/tests/components/vera/test_init.py +++ b/tests/components/vera/test_init.py @@ -13,6 +13,7 @@ from homeassistant.components.vera import ( ) from homeassistant.config_entries import ENTRY_STATE_NOT_LOADED from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er from .common import ComponentFactory, ConfigSource, new_simple_controller_config @@ -40,7 +41,7 @@ async def test_init( ), ) - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) entry1 = entity_registry.async_get(entity1_id) assert entry1 assert entry1.unique_id == "vera_first_serial_1" @@ -67,7 +68,7 @@ async def test_init_from_file( ), ) - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) entry1 = entity_registry.async_get(entity1_id) assert entry1 assert entry1.unique_id == "vera_first_serial_1" @@ -117,7 +118,7 @@ async def test_multiple_controllers_with_legacy_one( ), ) - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) entry1 = entity_registry.async_get(entity1_id) assert entry1 diff --git a/tests/components/wemo/conftest.py b/tests/components/wemo/conftest.py index 0e0a69216b2..69b4b84dcd3 100644 --- a/tests/components/wemo/conftest.py +++ b/tests/components/wemo/conftest.py @@ -7,6 +7,7 @@ import pywemo from homeassistant.components.wemo import CONF_DISCOVERY, CONF_STATIC from homeassistant.components.wemo.const import DOMAIN +from homeassistant.helpers import entity_registry as er from homeassistant.setup import async_setup_component MOCK_HOST = "127.0.0.1" @@ -72,7 +73,7 @@ async def async_wemo_entity_fixture(hass, pywemo_device): ) await hass.async_block_till_done() - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) entity_entries = list(entity_registry.entities.values()) assert len(entity_entries) == 1 diff --git a/tests/components/wemo/test_init.py b/tests/components/wemo/test_init.py index 374222d8688..1164af7cf95 100644 --- a/tests/components/wemo/test_init.py +++ b/tests/components/wemo/test_init.py @@ -6,6 +6,7 @@ import pywemo from homeassistant.components.wemo import CONF_DISCOVERY, CONF_STATIC, WemoDiscovery from homeassistant.components.wemo.const import DOMAIN +from homeassistant.helpers import entity_registry as er from homeassistant.setup import async_setup_component from homeassistant.util import dt @@ -41,7 +42,7 @@ async def test_static_duplicate_static_entry(hass, pywemo_device): }, ) await hass.async_block_till_done() - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) entity_entries = list(entity_reg.entities.values()) assert len(entity_entries) == 1 @@ -59,7 +60,7 @@ async def test_static_config_with_port(hass, pywemo_device): }, ) await hass.async_block_till_done() - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) entity_entries = list(entity_reg.entities.values()) assert len(entity_entries) == 1 @@ -77,7 +78,7 @@ async def test_static_config_without_port(hass, pywemo_device): }, ) await hass.async_block_till_done() - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) entity_entries = list(entity_reg.entities.values()) assert len(entity_entries) == 1 @@ -132,7 +133,7 @@ async def test_discovery(hass, pywemo_registry): await hass.async_block_till_done() # Verify that the expected number of devices were setup. - entity_reg = await hass.helpers.entity_registry.async_get_registry() + entity_reg = er.async_get(hass) entity_entries = list(entity_reg.entities.values()) assert len(entity_entries) == 3 diff --git a/tests/components/wilight/test_cover.py b/tests/components/wilight/test_cover.py index 85f62c9d120..8b058d95836 100644 --- a/tests/components/wilight/test_cover.py +++ b/tests/components/wilight/test_cover.py @@ -20,6 +20,7 @@ from homeassistant.const import ( STATE_OPEN, STATE_OPENING, ) +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.typing import HomeAssistantType from . import ( @@ -64,7 +65,7 @@ async def test_loading_cover( assert entry assert entry.unique_id == WILIGHT_ID - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) # First segment of the strip state = hass.states.get("cover.wl000000000099_1") diff --git a/tests/components/wilight/test_fan.py b/tests/components/wilight/test_fan.py index 1247b622ae7..0ad7789c52c 100644 --- a/tests/components/wilight/test_fan.py +++ b/tests/components/wilight/test_fan.py @@ -20,6 +20,7 @@ from homeassistant.const import ( STATE_OFF, STATE_ON, ) +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.typing import HomeAssistantType from . import ( @@ -64,7 +65,7 @@ async def test_loading_light_fan( assert entry assert entry.unique_id == WILIGHT_ID - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) # First segment of the strip state = hass.states.get("fan.wl000000000099_2") diff --git a/tests/components/wilight/test_light.py b/tests/components/wilight/test_light.py index b7250df546d..9abe17ce9e5 100644 --- a/tests/components/wilight/test_light.py +++ b/tests/components/wilight/test_light.py @@ -16,6 +16,7 @@ from homeassistant.const import ( STATE_OFF, STATE_ON, ) +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.typing import HomeAssistantType from tests.components.wilight import ( @@ -140,7 +141,7 @@ async def test_loading_light( assert entry assert entry.unique_id == WILIGHT_ID - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) # First segment of the strip state = hass.states.get("light.wl000000000099_1") diff --git a/tests/components/withings/test_binary_sensor.py b/tests/components/withings/test_binary_sensor.py index 3477671ea79..22c6b6de862 100644 --- a/tests/components/withings/test_binary_sensor.py +++ b/tests/components/withings/test_binary_sensor.py @@ -8,6 +8,7 @@ from homeassistant.components.withings.common import ( from homeassistant.components.withings.const import Measurement from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_registry import EntityRegistry from .common import ComponentFactory, new_profile_config @@ -21,9 +22,7 @@ async def test_binary_sensor( person0 = new_profile_config("person0", 0) person1 = new_profile_config("person1", 1) - entity_registry: EntityRegistry = ( - await hass.helpers.entity_registry.async_get_registry() - ) + entity_registry: EntityRegistry = er.async_get(hass) await component_factory.configure_component(profile_configs=(person0, person1)) assert not await async_get_entity_id(hass, in_bed_attribute, person0.user_id) diff --git a/tests/components/withings/test_sensor.py b/tests/components/withings/test_sensor.py index 16b83a585aa..8b51f62514d 100644 --- a/tests/components/withings/test_sensor.py +++ b/tests/components/withings/test_sensor.py @@ -27,6 +27,7 @@ from homeassistant.components.withings.common import ( ) from homeassistant.components.withings.const import Measurement from homeassistant.core import HomeAssistant, State +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_registry import EntityRegistry from .common import ComponentFactory, new_profile_config @@ -304,9 +305,7 @@ async def test_sensor_default_enabled_entities( hass: HomeAssistant, component_factory: ComponentFactory ) -> None: """Test entities enabled by default.""" - entity_registry: EntityRegistry = ( - await hass.helpers.entity_registry.async_get_registry() - ) + entity_registry: EntityRegistry = er.async_get(hass) await component_factory.configure_component(profile_configs=(PERSON0,)) @@ -347,9 +346,7 @@ async def test_all_entities( hass: HomeAssistant, component_factory: ComponentFactory ) -> None: """Test all entities.""" - entity_registry: EntityRegistry = ( - await hass.helpers.entity_registry.async_get_registry() - ) + entity_registry: EntityRegistry = er.async_get(hass) with patch( "homeassistant.components.withings.sensor.BaseWithingsSensor.entity_registry_enabled_default" diff --git a/tests/components/wled/test_light.py b/tests/components/wled/test_light.py index eb9124ab906..0077cea0202 100644 --- a/tests/components/wled/test_light.py +++ b/tests/components/wled/test_light.py @@ -36,6 +36,7 @@ from homeassistant.const import ( STATE_UNAVAILABLE, ) from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er import homeassistant.util.dt as dt_util from tests.common import async_fire_time_changed, load_fixture @@ -49,7 +50,7 @@ async def test_rgb_light_state( """Test the creation and values of the WLED lights.""" await init_integration(hass, aioclient_mock) - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) # First segment of the strip state = hass.states.get("light.wled_rgb_light_segment_0") diff --git a/tests/components/wled/test_sensor.py b/tests/components/wled/test_sensor.py index 11e14bd79d9..9cebf2cda32 100644 --- a/tests/components/wled/test_sensor.py +++ b/tests/components/wled/test_sensor.py @@ -23,6 +23,7 @@ from homeassistant.const import ( SIGNAL_STRENGTH_DECIBELS_MILLIWATT, ) from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er from homeassistant.util import dt as dt_util from tests.components.wled import init_integration @@ -35,7 +36,7 @@ async def test_sensors( """Test the creation and values of the WLED sensors.""" entry = await init_integration(hass, aioclient_mock, skip_setup=True) - registry = await hass.helpers.entity_registry.async_get_registry() + registry = er.async_get(hass) # Pre-create registry entries for disabled by default sensors registry.async_get_or_create( @@ -185,7 +186,7 @@ async def test_disabled_by_default_sensors( ) -> None: """Test the disabled by default WLED sensors.""" await init_integration(hass, aioclient_mock) - registry = await hass.helpers.entity_registry.async_get_registry() + registry = er.async_get(hass) state = hass.states.get(entity_id) assert state is None diff --git a/tests/components/wled/test_switch.py b/tests/components/wled/test_switch.py index c6e30ef903e..ddeeee41ac8 100644 --- a/tests/components/wled/test_switch.py +++ b/tests/components/wled/test_switch.py @@ -20,6 +20,7 @@ from homeassistant.const import ( STATE_UNAVAILABLE, ) from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er from tests.components.wled import init_integration from tests.test_util.aiohttp import AiohttpClientMocker @@ -31,7 +32,7 @@ async def test_switch_state( """Test the creation and values of the WLED switches.""" await init_integration(hass, aioclient_mock) - entity_registry = await hass.helpers.entity_registry.async_get_registry() + entity_registry = er.async_get(hass) state = hass.states.get("switch.wled_rgb_light_nightlight") assert state diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py index 05a0bd0d8d4..b6a59809d30 100644 --- a/tests/components/yeelight/test_init.py +++ b/tests/components/yeelight/test_init.py @@ -13,7 +13,7 @@ from homeassistant.components.yeelight import ( ) from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_NAME, STATE_UNAVAILABLE 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 . import ( @@ -106,11 +106,14 @@ async def test_unique_ids_device(hass: HomeAssistant): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() - er = await entity_registry.async_get_registry(hass) - assert er.async_get(ENTITY_BINARY_SENSOR).unique_id == f"{ID}-nightlight_sensor" - assert er.async_get(ENTITY_LIGHT).unique_id == ID - assert er.async_get(ENTITY_NIGHTLIGHT).unique_id == f"{ID}-nightlight" - assert er.async_get(ENTITY_AMBILIGHT).unique_id == f"{ID}-ambilight" + entity_registry = er.async_get(hass) + assert ( + entity_registry.async_get(ENTITY_BINARY_SENSOR).unique_id + == f"{ID}-nightlight_sensor" + ) + assert entity_registry.async_get(ENTITY_LIGHT).unique_id == ID + assert entity_registry.async_get(ENTITY_NIGHTLIGHT).unique_id == f"{ID}-nightlight" + assert entity_registry.async_get(ENTITY_AMBILIGHT).unique_id == f"{ID}-ambilight" async def test_unique_ids_entry(hass: HomeAssistant): @@ -131,18 +134,19 @@ async def test_unique_ids_entry(hass: HomeAssistant): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() - er = await entity_registry.async_get_registry(hass) + entity_registry = er.async_get(hass) assert ( - er.async_get(ENTITY_BINARY_SENSOR).unique_id + entity_registry.async_get(ENTITY_BINARY_SENSOR).unique_id == f"{config_entry.entry_id}-nightlight_sensor" ) - assert er.async_get(ENTITY_LIGHT).unique_id == config_entry.entry_id + assert entity_registry.async_get(ENTITY_LIGHT).unique_id == config_entry.entry_id assert ( - er.async_get(ENTITY_NIGHTLIGHT).unique_id + entity_registry.async_get(ENTITY_NIGHTLIGHT).unique_id == f"{config_entry.entry_id}-nightlight" ) assert ( - er.async_get(ENTITY_AMBILIGHT).unique_id == f"{config_entry.entry_id}-ambilight" + entity_registry.async_get(ENTITY_AMBILIGHT).unique_id + == f"{config_entry.entry_id}-ambilight" ) @@ -170,8 +174,8 @@ async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant): binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format( IP_ADDRESS.replace(".", "_") ) - er = await entity_registry.async_get_registry(hass) - assert er.async_get(binary_sensor_entity_id) is None + entity_registry = er.async_get(hass) + assert entity_registry.async_get(binary_sensor_entity_id) is None type(mocked_bulb).get_capabilities = MagicMock(CAPABILITIES) type(mocked_bulb).get_properties = MagicMock(None) @@ -179,5 +183,5 @@ async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant): hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE].update() await hass.async_block_till_done() - er = await entity_registry.async_get_registry(hass) - assert er.async_get(binary_sensor_entity_id) is not None + entity_registry = er.async_get(hass) + assert entity_registry.async_get(binary_sensor_entity_id) is not None diff --git a/tests/components/yeelight/test_light.py b/tests/components/yeelight/test_light.py index 0b7415140a3..90a8d1f9e6e 100644 --- a/tests/components/yeelight/test_light.py +++ b/tests/components/yeelight/test_light.py @@ -85,7 +85,7 @@ from homeassistant.components.yeelight.light import ( ) from homeassistant.const import ATTR_ENTITY_ID, CONF_HOST, CONF_NAME 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 homeassistant.util.color import ( color_hs_to_RGB, @@ -376,7 +376,7 @@ async def test_device_types(hass: HomeAssistant): await hass.config_entries.async_unload(config_entry.entry_id) await config_entry.async_remove(hass) - registry = await entity_registry.async_get_registry(hass) + registry = er.async_get(hass) registry.async_clear_config_entry(config_entry.entry_id) # nightlight diff --git a/tests/components/zha/test_device.py b/tests/components/zha/test_device.py index 1ce75045d38..d3ab3c3ada2 100644 --- a/tests/components/zha/test_device.py +++ b/tests/components/zha/test_device.py @@ -10,7 +10,7 @@ import zigpy.zcl.clusters.general as general import homeassistant.components.zha.core.device as zha_core_device from homeassistant.const import STATE_OFF, STATE_UNAVAILABLE -import homeassistant.helpers.device_registry as ha_dev_reg +import homeassistant.helpers.device_registry as dr import homeassistant.util.dt as dt_util from .common import async_enable_traffic, make_zcl_header @@ -227,7 +227,7 @@ async def test_ota_sw_version(hass, ota_zha_device): """Test device entry gets sw_version updated via OTA channel.""" ota_ch = ota_zha_device.channels.pools[0].client_channels["1:0x0019"] - dev_registry = await ha_dev_reg.async_get_registry(hass) + dev_registry = dr.async_get(hass) entry = dev_registry.async_get(ota_zha_device.device_id) assert entry.sw_version is None diff --git a/tests/components/zha/test_device_action.py b/tests/components/zha/test_device_action.py index 1160995e8d7..4a777fcebb6 100644 --- a/tests/components/zha/test_device_action.py +++ b/tests/components/zha/test_device_action.py @@ -12,7 +12,7 @@ from homeassistant.components.device_automation import ( _async_get_device_automations as async_get_device_automations, ) from homeassistant.components.zha import DOMAIN -from homeassistant.helpers.device_registry import async_get_registry +from homeassistant.helpers import device_registry as dr from homeassistant.setup import async_setup_component from tests.common import async_mock_service, mock_coro @@ -49,7 +49,7 @@ async def test_get_actions(hass, device_ias): ieee_address = str(device_ias[0].ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({(DOMAIN, ieee_address)}) actions = await async_get_device_automations(hass, "action", reg_device.id) @@ -72,7 +72,7 @@ async def test_action(hass, device_ias): ieee_address = str(zha_device.ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({(DOMAIN, ieee_address)}) with patch( diff --git a/tests/components/zha/test_device_trigger.py b/tests/components/zha/test_device_trigger.py index ec947846801..b2f964e2695 100644 --- a/tests/components/zha/test_device_trigger.py +++ b/tests/components/zha/test_device_trigger.py @@ -8,7 +8,7 @@ import zigpy.zcl.clusters.general as general import homeassistant.components.automation as automation import homeassistant.components.zha.core.device as zha_core_device -from homeassistant.helpers.device_registry import async_get_registry +from homeassistant.helpers import device_registry as dr from homeassistant.setup import async_setup_component import homeassistant.util.dt as dt_util @@ -86,7 +86,7 @@ async def test_triggers(hass, mock_devices): ieee_address = str(zha_device.ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({("zha", ieee_address)}) triggers = await async_get_device_automations(hass, "trigger", reg_device.id) @@ -144,7 +144,7 @@ async def test_no_triggers(hass, mock_devices): _, zha_device = mock_devices ieee_address = str(zha_device.ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({("zha", ieee_address)}) triggers = await async_get_device_automations(hass, "trigger", reg_device.id) @@ -173,7 +173,7 @@ async def test_if_fires_on_event(hass, mock_devices, calls): } ieee_address = str(zha_device.ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({("zha", ieee_address)}) assert await async_setup_component( @@ -282,7 +282,7 @@ async def test_exception_no_triggers(hass, mock_devices, calls, caplog): _, zha_device = mock_devices ieee_address = str(zha_device.ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({("zha", ieee_address)}) await async_setup_component( @@ -324,7 +324,7 @@ async def test_exception_bad_trigger(hass, mock_devices, calls, caplog): } ieee_address = str(zha_device.ieee) - ha_device_registry = await async_get_registry(hass) + ha_device_registry = dr.async_get(hass) reg_device = ha_device_registry.async_get_device({("zha", ieee_address)}) await async_setup_component( diff --git a/tests/components/zha/test_discover.py b/tests/components/zha/test_discover.py index d4195c681e7..c84c22e3251 100644 --- a/tests/components/zha/test_discover.py +++ b/tests/components/zha/test_discover.py @@ -67,7 +67,7 @@ async def test_devices( zha_device_joined_restored, ): """Test device discovery.""" - entity_registry = await homeassistant.helpers.entity_registry.async_get_registry( + entity_registry = homeassistant.helpers.entity_registry.async_get( hass_disable_services ) diff --git a/tests/components/zone/test_init.py b/tests/components/zone/test_init.py index 07fd83cbe77..38846050643 100644 --- a/tests/components/zone/test_init.py +++ b/tests/components/zone/test_init.py @@ -15,7 +15,7 @@ from homeassistant.const import ( ) from homeassistant.core import Context from homeassistant.exceptions import Unauthorized -from homeassistant.helpers import entity_registry +from homeassistant.helpers import entity_registry as er from tests.common import MockConfigEntry @@ -244,7 +244,7 @@ async def test_core_config_update(hass): async def test_reload(hass, hass_admin_user, hass_read_only_user): """Test reload service.""" count_start = len(hass.states.async_entity_ids()) - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) assert await setup.async_setup_component( hass, @@ -365,7 +365,7 @@ async def test_ws_delete(hass, hass_ws_client, storage_setup): input_id = "from_storage" input_entity_id = f"{DOMAIN}.{input_id}" - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) state = hass.states.get(input_entity_id) assert state is not None @@ -401,7 +401,7 @@ async def test_update(hass, hass_ws_client, storage_setup): input_id = "from_storage" input_entity_id = f"{DOMAIN}.{input_id}" - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) state = hass.states.get(input_entity_id) assert state.attributes["latitude"] == 1 @@ -435,7 +435,7 @@ async def test_ws_create(hass, hass_ws_client, storage_setup): input_id = "new_input" input_entity_id = f"{DOMAIN}.{input_id}" - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) state = hass.states.get(input_entity_id) assert state is None diff --git a/tests/components/zwave/test_init.py b/tests/components/zwave/test_init.py index d64e191e1f3..878513da579 100644 --- a/tests/components/zwave/test_init.py +++ b/tests/components/zwave/test_init.py @@ -18,8 +18,7 @@ from homeassistant.components.zwave import ( ) from homeassistant.components.zwave.binary_sensor import get_device from homeassistant.const import ATTR_ENTITY_ID, ATTR_NAME -from homeassistant.helpers.device_registry import async_get_registry as get_dev_reg -from homeassistant.helpers.entity_registry import async_get_registry +from homeassistant.helpers import device_registry as dr, entity_registry as er from tests.common import async_fire_time_changed, mock_registry from tests.mock.zwave import MockEntityValues, MockNetwork, MockNode, MockValue @@ -472,8 +471,8 @@ async def test_value_entities(hass, mock_openzwave): assert hass.states.get("binary_sensor.mock_node_mock_value").state == "off" assert hass.states.get("binary_sensor.mock_node_mock_value_b").state == "off" - ent_reg = await async_get_registry(hass) - dev_reg = await get_dev_reg(hass) + ent_reg = er.async_get(hass) + dev_reg = dr.async_get(hass) entry = ent_reg.async_get("zwave.mock_node") assert entry is not None diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py index 78bac5518a0..aa085836b65 100644 --- a/tests/components/zwave_js/test_api.py +++ b/tests/components/zwave_js/test_api.py @@ -23,7 +23,7 @@ from homeassistant.components.zwave_js.api import ( VALUE, ) from homeassistant.components.zwave_js.const import DOMAIN -from homeassistant.helpers.device_registry import async_get_registry +from homeassistant.helpers import device_registry as dr async def test_websocket_api(hass, integration, multisensor_6, hass_ws_client): @@ -197,7 +197,7 @@ async def test_remove_node( # Add mock node to controller client.driver.controller.nodes[67] = nortek_thermostat - dev_reg = await async_get_registry(hass) + dev_reg = dr.async_get(hass) # Create device registry entry for mock node device = dev_reg.async_get_or_create( diff --git a/tests/components/zwave_js/test_binary_sensor.py b/tests/components/zwave_js/test_binary_sensor.py index e8361d8b03e..ddfed9727e6 100644 --- a/tests/components/zwave_js/test_binary_sensor.py +++ b/tests/components/zwave_js/test_binary_sensor.py @@ -3,6 +3,7 @@ from zwave_js_server.event import Event from homeassistant.components.binary_sensor import DEVICE_CLASS_MOTION from homeassistant.const import DEVICE_CLASS_BATTERY, STATE_OFF, STATE_ON +from homeassistant.helpers import entity_registry as er from .common import ( DISABLED_LEGACY_BINARY_SENSOR, @@ -61,7 +62,7 @@ async def test_disabled_legacy_sensor(hass, multisensor_6, integration): """Test disabled legacy boolean binary sensor.""" # this node has Notification CC implemented so legacy binary sensor should be disabled - registry = await hass.helpers.entity_registry.async_get_registry() + registry = er.async_get(hass) entity_id = DISABLED_LEGACY_BINARY_SENSOR state = hass.states.get(entity_id) assert state is None diff --git a/tests/components/zwave_js/test_init.py b/tests/components/zwave_js/test_init.py index 00574bd2d2f..3e7f79b9cec 100644 --- a/tests/components/zwave_js/test_init.py +++ b/tests/components/zwave_js/test_init.py @@ -17,7 +17,7 @@ from homeassistant.config_entries import ( ENTRY_STATE_SETUP_RETRY, ) from homeassistant.const import STATE_UNAVAILABLE -from homeassistant.helpers import device_registry, entity_registry +from homeassistant.helpers import device_registry as dr, entity_registry as er from .common import ( AIR_TEMPERATURE_SENSOR, @@ -117,7 +117,7 @@ async def test_unique_id_migration_dupes( hass, multisensor_6_state, client, integration ): """Test we remove an entity when .""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) entity_name = AIR_TEMPERATURE_SENSOR.split(".")[1] @@ -168,7 +168,7 @@ async def test_unique_id_migration_dupes( async def test_unique_id_migration_v1(hass, multisensor_6_state, client, integration): """Test unique ID is migrated from old format to new (version 1).""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) # Migrate version 1 entity_name = AIR_TEMPERATURE_SENSOR.split(".")[1] @@ -201,7 +201,7 @@ async def test_unique_id_migration_v1(hass, multisensor_6_state, client, integra async def test_unique_id_migration_v2(hass, multisensor_6_state, client, integration): """Test unique ID is migrated from old format to new (version 2).""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) # Migrate version 2 ILLUMINANCE_SENSOR = "sensor.multisensor_6_illuminance" entity_name = ILLUMINANCE_SENSOR.split(".")[1] @@ -234,7 +234,7 @@ async def test_unique_id_migration_v2(hass, multisensor_6_state, client, integra async def test_unique_id_migration_v3(hass, multisensor_6_state, client, integration): """Test unique ID is migrated from old format to new (version 3).""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) # Migrate version 2 ILLUMINANCE_SENSOR = "sensor.multisensor_6_illuminance" entity_name = ILLUMINANCE_SENSOR.split(".")[1] @@ -269,7 +269,7 @@ async def test_unique_id_migration_property_key_v1( hass, hank_binary_switch_state, client, integration ): """Test unique ID with property key is migrated from old format to new (version 1).""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) SENSOR_NAME = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed" entity_name = SENSOR_NAME.split(".")[1] @@ -304,7 +304,7 @@ async def test_unique_id_migration_property_key_v2( hass, hank_binary_switch_state, client, integration ): """Test unique ID with property key is migrated from old format to new (version 2).""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) SENSOR_NAME = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed" entity_name = SENSOR_NAME.split(".")[1] @@ -341,7 +341,7 @@ async def test_unique_id_migration_property_key_v3( hass, hank_binary_switch_state, client, integration ): """Test unique ID with property key is migrated from old format to new (version 3).""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) SENSOR_NAME = "sensor.smart_plug_with_two_usb_ports_value_electric_consumed" entity_name = SENSOR_NAME.split(".")[1] @@ -376,7 +376,7 @@ async def test_unique_id_migration_notification_binary_sensor( hass, multisensor_6_state, client, integration ): """Test unique ID is migrated from old format to new for a notification binary sensor.""" - ent_reg = entity_registry.async_get(hass) + ent_reg = er.async_get(hass) entity_name = NOTIFICATION_MOTION_BINARY_SENSOR.split(".")[1] @@ -813,17 +813,13 @@ async def test_removed_device(hass, client, multiple_devices, integration): assert len(client.driver.controller.nodes) == 2 # Make sure there are the same number of devices - dev_reg = await device_registry.async_get_registry(hass) - device_entries = device_registry.async_entries_for_config_entry( - dev_reg, integration.entry_id - ) + dev_reg = dr.async_get(hass) + device_entries = dr.async_entries_for_config_entry(dev_reg, integration.entry_id) assert len(device_entries) == 2 # Check how many entities there are - ent_reg = await entity_registry.async_get_registry(hass) - entity_entries = entity_registry.async_entries_for_config_entry( - ent_reg, integration.entry_id - ) + ent_reg = er.async_get(hass) + entity_entries = er.async_entries_for_config_entry(ent_reg, integration.entry_id) assert len(entity_entries) == 24 # Remove a node and reload the entry @@ -833,21 +829,17 @@ async def test_removed_device(hass, client, multiple_devices, integration): # Assert that the node and all of it's entities were removed from the device and # entity registry - device_entries = device_registry.async_entries_for_config_entry( - dev_reg, integration.entry_id - ) + device_entries = dr.async_entries_for_config_entry(dev_reg, integration.entry_id) assert len(device_entries) == 1 - entity_entries = entity_registry.async_entries_for_config_entry( - ent_reg, integration.entry_id - ) + entity_entries = er.async_entries_for_config_entry(ent_reg, integration.entry_id) assert len(entity_entries) == 15 assert dev_reg.async_get_device({get_device_id(client, old_node)}) is None async def test_suggested_area(hass, client, eaton_rf9640_dimmer): """Test that suggested area works.""" - dev_reg = device_registry.async_get(hass) - ent_reg = entity_registry.async_get(hass) + dev_reg = dr.async_get(hass) + ent_reg = er.async_get(hass) entry = MockConfigEntry(domain="zwave_js", data={"url": "ws://test.org"}) entry.add_to_hass(hass) diff --git a/tests/components/zwave_js/test_sensor.py b/tests/components/zwave_js/test_sensor.py index bd6fb9f2569..398d831e692 100644 --- a/tests/components/zwave_js/test_sensor.py +++ b/tests/components/zwave_js/test_sensor.py @@ -7,10 +7,7 @@ from homeassistant.const import ( POWER_WATT, TEMP_CELSIUS, ) -from homeassistant.helpers.entity_registry import ( - DISABLED_INTEGRATION, - async_get_registry, -) +from homeassistant.helpers import entity_registry as er from .common import ( AIR_TEMPERATURE_SENSOR, @@ -49,12 +46,12 @@ async def test_energy_sensors(hass, hank_binary_switch, integration): async def test_disabled_notification_sensor(hass, multisensor_6, integration): """Test sensor is created from Notification CC and is disabled.""" - ent_reg = await async_get_registry(hass) + ent_reg = er.async_get(hass) entity_entry = ent_reg.async_get(NOTIFICATION_MOTION_SENSOR) assert entity_entry assert entity_entry.disabled - assert entity_entry.disabled_by == DISABLED_INTEGRATION + assert entity_entry.disabled_by == er.DISABLED_INTEGRATION # Test enabling entity updated_entry = ent_reg.async_update_entity( diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index 965ebcd3e23..1d1b1a3cca3 100644 --- a/tests/helpers/test_device_registry.py +++ b/tests/helpers/test_device_registry.py @@ -797,7 +797,7 @@ async def test_cleanup_device_registry(hass, registry): identifiers={("something", "d4")}, config_entry_id="non_existing" ) - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = entity_registry.async_get(hass) ent_reg.async_get_or_create("light", "hue", "e1", device_id=d1.id) ent_reg.async_get_or_create("light", "hue", "e2", device_id=d1.id) ent_reg.async_get_or_create("light", "hue", "e3", device_id=d3.id) @@ -829,7 +829,7 @@ async def test_cleanup_device_registry_removes_expired_orphaned_devices(hass, re assert len(registry.devices) == 0 assert len(registry.deleted_devices) == 3 - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = entity_registry.async_get(hass) device_registry.async_cleanup(hass, registry, ent_reg) assert len(registry.devices) == 0 @@ -847,7 +847,6 @@ async def test_cleanup_device_registry_removes_expired_orphaned_devices(hass, re async def test_cleanup_startup(hass): """Test we run a cleanup on startup.""" hass.state = CoreState.not_running - await device_registry.async_get_registry(hass) with patch( "homeassistant.helpers.device_registry.Debouncer.async_call" diff --git a/tests/helpers/test_entity_platform.py b/tests/helpers/test_entity_platform.py index ab3e04843f9..3f26535de18 100644 --- a/tests/helpers/test_entity_platform.py +++ b/tests/helpers/test_entity_platform.py @@ -9,7 +9,11 @@ import pytest from homeassistant.const import PERCENTAGE from homeassistant.core import callback from homeassistant.exceptions import HomeAssistantError, PlatformNotReady -from homeassistant.helpers import entity_platform, entity_registry +from homeassistant.helpers import ( + device_registry as dr, + entity_platform, + entity_registry as er, +) from homeassistant.helpers.entity import async_generate_entity_id from homeassistant.helpers.entity_component import ( DEFAULT_SCAN_INTERVAL, @@ -467,7 +471,7 @@ async def test_overriding_name_from_registry(hass): mock_registry( hass, { - "test_domain.world": entity_registry.RegistryEntry( + "test_domain.world": er.RegistryEntry( entity_id="test_domain.world", unique_id="1234", # Using component.async_add_entities is equal to platform "domain" @@ -499,12 +503,12 @@ async def test_registry_respect_entity_disabled(hass): mock_registry( hass, { - "test_domain.world": entity_registry.RegistryEntry( + "test_domain.world": er.RegistryEntry( entity_id="test_domain.world", unique_id="1234", # Using component.async_add_entities is equal to platform "domain" platform="test_platform", - disabled_by=entity_registry.DISABLED_USER, + disabled_by=er.DISABLED_USER, ) }, ) @@ -520,7 +524,7 @@ async def test_entity_registry_updates_name(hass): registry = mock_registry( hass, { - "test_domain.world": entity_registry.RegistryEntry( + "test_domain.world": er.RegistryEntry( entity_id="test_domain.world", unique_id="1234", # Using component.async_add_entities is equal to platform "domain" @@ -624,7 +628,7 @@ async def test_entity_registry_updates_entity_id(hass): registry = mock_registry( hass, { - "test_domain.world": entity_registry.RegistryEntry( + "test_domain.world": er.RegistryEntry( entity_id="test_domain.world", unique_id="1234", # Using component.async_add_entities is equal to platform "domain" @@ -656,14 +660,14 @@ async def test_entity_registry_updates_invalid_entity_id(hass): registry = mock_registry( hass, { - "test_domain.world": entity_registry.RegistryEntry( + "test_domain.world": er.RegistryEntry( entity_id="test_domain.world", unique_id="1234", # Using component.async_add_entities is equal to platform "domain" platform="test_platform", name="Some name", ), - "test_domain.existing": entity_registry.RegistryEntry( + "test_domain.existing": er.RegistryEntry( entity_id="test_domain.existing", unique_id="5678", platform="test_platform", @@ -703,7 +707,7 @@ async def test_entity_registry_updates_invalid_entity_id(hass): async def test_device_info_called(hass): """Test device info is forwarded correctly.""" - registry = await hass.helpers.device_registry.async_get_registry() + registry = dr.async_get(hass) via = registry.async_get_or_create( config_entry_id="123", connections=set(), @@ -763,7 +767,7 @@ async def test_device_info_called(hass): async def test_device_info_not_overrides(hass): """Test device info is forwarded correctly.""" - registry = await hass.helpers.device_registry.async_get_registry() + registry = dr.async_get(hass) device = registry.async_get_or_create( config_entry_id="bla", connections={("mac", "abcd")}, @@ -823,7 +827,7 @@ async def test_entity_disabled_by_integration(hass): assert entity_disabled.hass is None assert entity_disabled.platform is None - registry = await hass.helpers.entity_registry.async_get_registry() + registry = er.async_get(hass) entry_default = registry.async_get_or_create(DOMAIN, DOMAIN, "default") assert entry_default.disabled_by is None @@ -845,7 +849,7 @@ async def test_entity_info_added_to_entity_registry(hass): await component.async_add_entities([entity_default]) - registry = await hass.helpers.entity_registry.async_get_registry() + registry = er.async_get(hass) entry_default = registry.async_get_or_create(DOMAIN, DOMAIN, "default") print(entry_default) diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index 0a1a27efef5..d671aacebb3 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -5,7 +5,7 @@ import pytest from homeassistant.const import EVENT_HOMEASSISTANT_START, STATE_UNAVAILABLE from homeassistant.core import CoreState, callback, valid_entity_id -from homeassistant.helpers import entity_registry +from homeassistant.helpers import entity_registry as er from tests.common import ( MockConfigEntry, @@ -32,7 +32,7 @@ def update_events(hass): def async_capture(event): events.append(event.data) - hass.bus.async_listen(entity_registry.EVENT_ENTITY_REGISTRY_UPDATED, async_capture) + hass.bus.async_listen(er.EVENT_ENTITY_REGISTRY_UPDATED, async_capture) return events @@ -74,7 +74,7 @@ def test_get_or_create_updates_data(registry): capabilities={"max": 100}, supported_features=5, device_class="mock-device-class", - disabled_by=entity_registry.DISABLED_HASS, + disabled_by=er.DISABLED_HASS, unit_of_measurement="initial-unit_of_measurement", original_name="initial-original_name", original_icon="initial-original_icon", @@ -85,7 +85,7 @@ def test_get_or_create_updates_data(registry): assert orig_entry.capabilities == {"max": 100} assert orig_entry.supported_features == 5 assert orig_entry.device_class == "mock-device-class" - assert orig_entry.disabled_by == entity_registry.DISABLED_HASS + assert orig_entry.disabled_by == er.DISABLED_HASS assert orig_entry.unit_of_measurement == "initial-unit_of_measurement" assert orig_entry.original_name == "initial-original_name" assert orig_entry.original_icon == "initial-original_icon" @@ -101,7 +101,7 @@ def test_get_or_create_updates_data(registry): capabilities={"new-max": 100}, supported_features=10, device_class="new-mock-device-class", - disabled_by=entity_registry.DISABLED_USER, + disabled_by=er.DISABLED_USER, unit_of_measurement="updated-unit_of_measurement", original_name="updated-original_name", original_icon="updated-original_icon", @@ -116,7 +116,7 @@ def test_get_or_create_updates_data(registry): assert new_entry.original_name == "updated-original_name" assert new_entry.original_icon == "updated-original_icon" # Should not be updated - assert new_entry.disabled_by == entity_registry.DISABLED_HASS + assert new_entry.disabled_by == er.DISABLED_HASS def test_get_or_create_suggested_object_id_conflict_register(registry): @@ -162,7 +162,7 @@ async def test_loading_saving_data(hass, registry): capabilities={"max": 100}, supported_features=5, device_class="mock-device-class", - disabled_by=entity_registry.DISABLED_HASS, + disabled_by=er.DISABLED_HASS, original_name="Original Name", original_icon="hass:original-icon", ) @@ -173,7 +173,7 @@ async def test_loading_saving_data(hass, registry): assert len(registry.entities) == 2 # Now load written data in new registry - registry2 = entity_registry.EntityRegistry(hass) + registry2 = er.EntityRegistry(hass) await flush_store(registry._store) await registry2.async_load() @@ -187,7 +187,7 @@ async def test_loading_saving_data(hass, registry): assert new_entry2.device_id == "mock-dev-id" assert new_entry2.area_id == "mock-area-id" - assert new_entry2.disabled_by == entity_registry.DISABLED_HASS + assert new_entry2.disabled_by == er.DISABLED_HASS assert new_entry2.capabilities == {"max": 100} assert new_entry2.supported_features == 5 assert new_entry2.device_class == "mock-device-class" @@ -220,8 +220,8 @@ def test_is_registered(registry): @pytest.mark.parametrize("load_registries", [False]) async def test_loading_extra_values(hass, hass_storage): """Test we load extra data from the registry.""" - hass_storage[entity_registry.STORAGE_KEY] = { - "version": entity_registry.STORAGE_VERSION, + hass_storage[er.STORAGE_KEY] = { + "version": er.STORAGE_VERSION, "data": { "entities": [ { @@ -257,8 +257,8 @@ async def test_loading_extra_values(hass, hass_storage): }, } - await entity_registry.async_load(hass) - registry = entity_registry.async_get(hass) + await er.async_load(hass) + registry = er.async_get(hass) assert len(registry.entities) == 4 @@ -279,9 +279,9 @@ async def test_loading_extra_values(hass, hass_storage): "test", "super_platform", "disabled-user" ) assert entry_disabled_hass.disabled - assert entry_disabled_hass.disabled_by == entity_registry.DISABLED_HASS + assert entry_disabled_hass.disabled_by == er.DISABLED_HASS assert entry_disabled_user.disabled - assert entry_disabled_user.disabled_by == entity_registry.DISABLED_USER + assert entry_disabled_user.disabled_by == er.DISABLED_USER def test_async_get_entity_id(registry): @@ -367,8 +367,8 @@ async def test_migration(hass): with patch("os.path.isfile", return_value=True), patch("os.remove"), patch( "homeassistant.helpers.entity_registry.load_yaml", return_value=old_conf ): - await entity_registry.async_load(hass) - registry = entity_registry.async_get(hass) + await er.async_load(hass) + registry = er.async_get(hass) assert registry.async_is_registered("light.kitchen") entry = registry.async_get_or_create( @@ -384,8 +384,8 @@ async def test_migration(hass): async def test_loading_invalid_entity_id(hass, hass_storage): """Test we autofix invalid entity IDs.""" - hass_storage[entity_registry.STORAGE_KEY] = { - "version": entity_registry.STORAGE_VERSION, + hass_storage[er.STORAGE_KEY] = { + "version": er.STORAGE_VERSION, "data": { "entities": [ { @@ -408,7 +408,7 @@ async def test_loading_invalid_entity_id(hass, hass_storage): }, } - registry = await entity_registry.async_get_registry(hass) + registry = er.async_get(hass) entity_invalid_middle = registry.async_get_or_create( "test", "super_platform", "id-invalid-middle" @@ -479,7 +479,7 @@ async def test_update_entity(registry): for attr_name, new_value in ( ("name", "new name"), ("icon", "new icon"), - ("disabled_by", entity_registry.DISABLED_USER), + ("disabled_by", er.DISABLED_USER), ): changes = {attr_name: new_value} updated_entry = registry.async_update_entity(entry.entity_id, **changes) @@ -531,7 +531,7 @@ async def test_restore_states(hass): """Test restoring states.""" hass.state = CoreState.not_running - registry = await entity_registry.async_get_registry(hass) + registry = er.async_get(hass) registry.async_get_or_create( "light", @@ -545,7 +545,7 @@ async def test_restore_states(hass): "hue", "5678", suggested_object_id="disabled", - disabled_by=entity_registry.DISABLED_HASS, + disabled_by=er.DISABLED_HASS, ) registry.async_get_or_create( "light", @@ -597,7 +597,7 @@ async def test_async_get_device_class_lookup(hass): """Test registry device class lookup.""" hass.state = CoreState.not_running - ent_reg = await entity_registry.async_get_registry(hass) + ent_reg = er.async_get(hass) ent_reg.async_get_or_create( "binary_sensor", @@ -888,10 +888,10 @@ async def test_disabled_entities_excluded_from_entity_list(hass, registry): disabled_by="user", ) - entries = entity_registry.async_entries_for_device(registry, device_entry.id) + entries = er.async_entries_for_device(registry, device_entry.id) assert entries == [entry1] - entries = entity_registry.async_entries_for_device( + entries = er.async_entries_for_device( registry, device_entry.id, include_disabled_entities=True ) assert entries == [entry1, entry2]