From d5797d9f7dabbfa2ecb7685b82cad036cbfba2b3 Mon Sep 17 00:00:00 2001
From: Shay Levy <levyshay1@gmail.com>
Date: Sun, 22 Jan 2023 01:26:54 +0200
Subject: [PATCH] Cleanup Shelly imports (#86359)

* Cleanup Shelly imports

* Cleanup tests
---
 homeassistant/components/shelly/__init__.py   | 35 ++++------
 .../components/shelly/bluetooth/scanner.py    |  7 +-
 homeassistant/components/shelly/climate.py    | 21 +++---
 .../components/shelly/config_flow.py          | 55 +++++++--------
 .../components/shelly/coordinator.py          | 22 +++---
 homeassistant/components/shelly/entity.py     | 36 +++++-----
 homeassistant/components/shelly/utils.py      | 25 ++++---
 tests/components/shelly/__init__.py           |  9 +--
 tests/components/shelly/test_config_flow.py   | 68 +++++++++++--------
 .../components/shelly/test_device_trigger.py  |  4 +-
 tests/components/shelly/test_init.py          | 16 +----
 11 files changed, 141 insertions(+), 157 deletions(-)

diff --git a/homeassistant/components/shelly/__init__.py b/homeassistant/components/shelly/__init__.py
index df054598f5c..f75f220d435 100644
--- a/homeassistant/components/shelly/__init__.py
+++ b/homeassistant/components/shelly/__init__.py
@@ -4,8 +4,8 @@ from __future__ import annotations
 import contextlib
 from typing import Any, Final
 
-import aioshelly
 from aioshelly.block_device import BlockDevice
+from aioshelly.common import ConnectionOptions
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
 from aioshelly.rpc_device import RpcDevice, UpdateType
 import voluptuous as vol
@@ -14,8 +14,13 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import aiohttp_client, device_registry
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.device_registry import (
+    CONNECTION_NETWORK_MAC,
+    async_get as dr_async_get,
+    format_mac,
+)
 from homeassistant.helpers.typing import ConfigType
 
 from .const import (
@@ -116,7 +121,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Shelly block based device from a config entry."""
-    options = aioshelly.common.ConnectionOptions(
+    options = ConnectionOptions(
         entry.data[CONF_HOST],
         entry.data.get(CONF_USERNAME),
         entry.data.get(CONF_PASSWORD),
@@ -125,23 +130,18 @@ async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> b
     coap_context = await get_coap_context(hass)
 
     device = await BlockDevice.create(
-        aiohttp_client.async_get_clientsession(hass),
+        async_get_clientsession(hass),
         coap_context,
         options,
         False,
     )
 
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr_async_get(hass)
     device_entry = None
     if entry.unique_id is not None:
         device_entry = dev_reg.async_get_device(
             identifiers=set(),
-            connections={
-                (
-                    device_registry.CONNECTION_NETWORK_MAC,
-                    device_registry.format_mac(entry.unique_id),
-                )
-            },
+            connections={(CONNECTION_NETWORK_MAC, format_mac(entry.unique_id))},
         )
     # https://github.com/home-assistant/core/pull/48076
     if device_entry and entry.entry_id not in device_entry.config_entries:
@@ -205,7 +205,7 @@ async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> b
 
 async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Shelly RPC based device from a config entry."""
-    options = aioshelly.common.ConnectionOptions(
+    options = ConnectionOptions(
         entry.data[CONF_HOST],
         entry.data.get(CONF_USERNAME),
         entry.data.get(CONF_PASSWORD),
@@ -214,23 +214,18 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> boo
     ws_context = await get_ws_context(hass)
 
     device = await RpcDevice.create(
-        aiohttp_client.async_get_clientsession(hass),
+        async_get_clientsession(hass),
         ws_context,
         options,
         False,
     )
 
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr_async_get(hass)
     device_entry = None
     if entry.unique_id is not None:
         device_entry = dev_reg.async_get_device(
             identifiers=set(),
-            connections={
-                (
-                    device_registry.CONNECTION_NETWORK_MAC,
-                    device_registry.format_mac(entry.unique_id),
-                )
-            },
+            connections={(CONNECTION_NETWORK_MAC, format_mac(entry.unique_id))},
         )
     # https://github.com/home-assistant/core/pull/48076
     if device_entry and entry.entry_id not in device_entry.config_entries:
diff --git a/homeassistant/components/shelly/bluetooth/scanner.py b/homeassistant/components/shelly/bluetooth/scanner.py
index f255d01c78b..5b302e0da62 100644
--- a/homeassistant/components/shelly/bluetooth/scanner.py
+++ b/homeassistant/components/shelly/bluetooth/scanner.py
@@ -1,7 +1,6 @@
 """Bluetooth scanner for shelly."""
 from __future__ import annotations
 
-import logging
 from typing import Any
 
 from aioshelly.ble import parse_ble_scan_result_event
@@ -10,7 +9,7 @@ from aioshelly.ble.const import BLE_SCAN_RESULT_EVENT, BLE_SCAN_RESULT_VERSION
 from homeassistant.components.bluetooth import BaseHaRemoteScanner
 from homeassistant.core import callback
 
-_LOGGER = logging.getLogger(__name__)
+from ..const import LOGGER
 
 
 class ShellyBLEScanner(BaseHaRemoteScanner):
@@ -25,7 +24,7 @@ class ShellyBLEScanner(BaseHaRemoteScanner):
         data = event["data"]
 
         if data[0] != BLE_SCAN_RESULT_VERSION:
-            _LOGGER.warning("Unsupported BLE scan result version: %s", data[0])
+            LOGGER.warning("Unsupported BLE scan result version: %s", data[0])
             return
 
         try:
@@ -33,7 +32,7 @@ class ShellyBLEScanner(BaseHaRemoteScanner):
         except Exception as err:  # pylint: disable=broad-except
             # Broad exception catch because we have no
             # control over the data that is coming in.
-            _LOGGER.error("Failed to parse BLE event: %s", err, exc_info=True)
+            LOGGER.error("Failed to parse BLE event: %s", err, exc_info=True)
             return
 
         self._async_on_advertisement(
diff --git a/homeassistant/components/shelly/climate.py b/homeassistant/components/shelly/climate.py
index b9435a56ebd..a2b81bc222c 100644
--- a/homeassistant/components/shelly/climate.py
+++ b/homeassistant/components/shelly/climate.py
@@ -19,9 +19,14 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant, State, callback
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import device_registry, entity_registry
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.entity_registry import (
+    RegistryEntry,
+    async_entries_for_config_entry,
+    async_get as er_async_get,
+)
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util.unit_conversion import TemperatureConverter
@@ -81,10 +86,8 @@ def async_restore_climate_entities(
 ) -> None:
     """Restore sleeping climate devices."""
 
-    ent_reg = entity_registry.async_get(hass)
-    entries = entity_registry.async_entries_for_config_entry(
-        ent_reg, config_entry.entry_id
-    )
+    ent_reg = er_async_get(hass)
+    entries = async_entries_for_config_entry(ent_reg, config_entry.entry_id)
 
     for entry in entries:
 
@@ -117,7 +120,7 @@ class BlockSleepingClimate(
         coordinator: ShellyBlockCoordinator,
         sensor_block: Block | None,
         device_block: Block | None,
-        entry: entity_registry.RegistryEntry | None = None,
+        entry: RegistryEntry | None = None,
     ) -> None:
         """Initialize climate."""
         super().__init__(coordinator)
@@ -242,11 +245,7 @@ class BlockSleepingClimate(
     @property
     def device_info(self) -> DeviceInfo:
         """Device info."""
-        return {
-            "connections": {
-                (device_registry.CONNECTION_NETWORK_MAC, self.coordinator.mac)
-            }
-        }
+        return {"connections": {(CONNECTION_NETWORK_MAC, self.coordinator.mac)}}
 
     def _check_is_off(self) -> bool:
         """Return if valve is off or on."""
diff --git a/homeassistant/components/shelly/config_flow.py b/homeassistant/components/shelly/config_flow.py
index 3b24bf026a9..612e2be0a74 100644
--- a/homeassistant/components/shelly/config_flow.py
+++ b/homeassistant/components/shelly/config_flow.py
@@ -4,8 +4,8 @@ from __future__ import annotations
 from collections.abc import Mapping
 from typing import Any, Final
 
-import aioshelly
 from aioshelly.block_device import BlockDevice
+from aioshelly.common import ConnectionOptions, get_info
 from aioshelly.exceptions import (
     DeviceConnectionError,
     FirmwareUnsupported,
@@ -15,12 +15,17 @@ from aioshelly.rpc_device import RpcDevice
 from awesomeversion import AwesomeVersion
 import voluptuous as vol
 
-from homeassistant import config_entries
-from homeassistant.components import zeroconf
+from homeassistant.components.zeroconf import ZeroconfServiceInfo
+from homeassistant.config_entries import ConfigEntry, ConfigFlow, OptionsFlow
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.data_entry_flow import FlowResult
-from homeassistant.helpers import aiohttp_client, selector
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.selector import (
+    SelectOptionDict,
+    SelectSelector,
+    SelectSelectorConfig,
+)
 
 from .const import (
     BLE_MIN_VERSION,
@@ -48,9 +53,9 @@ HOST_SCHEMA: Final = vol.Schema({vol.Required(CONF_HOST): str})
 
 
 BLE_SCANNER_OPTIONS = [
-    selector.SelectOptionDict(value=BLEScannerMode.DISABLED, label="Disabled"),
-    selector.SelectOptionDict(value=BLEScannerMode.ACTIVE, label="Active"),
-    selector.SelectOptionDict(value=BLEScannerMode.PASSIVE, label="Passive"),
+    SelectOptionDict(value=BLEScannerMode.DISABLED, label="Disabled"),
+    SelectOptionDict(value=BLEScannerMode.ACTIVE, label="Active"),
+    SelectOptionDict(value=BLEScannerMode.PASSIVE, label="Passive"),
 ]
 
 INTERNAL_WIFI_AP_IP = "192.168.33.1"
@@ -66,16 +71,12 @@ async def validate_input(
 
     Data has the keys from HOST_SCHEMA with values provided by the user.
     """
-    options = aioshelly.common.ConnectionOptions(
-        host,
-        data.get(CONF_USERNAME),
-        data.get(CONF_PASSWORD),
-    )
+    options = ConnectionOptions(host, data.get(CONF_USERNAME), data.get(CONF_PASSWORD))
 
     if get_info_gen(info) == 2:
         ws_context = await get_ws_context(hass)
         rpc_device = await RpcDevice.create(
-            aiohttp_client.async_get_clientsession(hass),
+            async_get_clientsession(hass),
             ws_context,
             options,
         )
@@ -92,7 +93,7 @@ async def validate_input(
     # Gen1
     coap_context = await get_coap_context(hass)
     block_device = await BlockDevice.create(
-        aiohttp_client.async_get_clientsession(hass),
+        async_get_clientsession(hass),
         coap_context,
         options,
     )
@@ -105,7 +106,7 @@ async def validate_input(
     }
 
 
-class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+class ShellyConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Shelly."""
 
     VERSION = 1
@@ -113,7 +114,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     host: str = ""
     info: dict[str, Any] = {}
     device_info: dict[str, Any] = {}
-    entry: config_entries.ConfigEntry | None = None
+    entry: ConfigEntry | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -232,7 +233,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             self._abort_if_unique_id_configured({CONF_HOST: host})
 
     async def async_step_zeroconf(
-        self, discovery_info: zeroconf.ZeroconfServiceInfo
+        self, discovery_info: ZeroconfServiceInfo
     ) -> FlowResult:
         """Handle zeroconf discovery."""
         host = discovery_info.host
@@ -352,33 +353,27 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def _async_get_info(self, host: str) -> dict[str, Any]:
         """Get info from shelly device."""
-        return await aioshelly.common.get_info(
-            aiohttp_client.async_get_clientsession(self.hass), host
-        )
+        return await get_info(async_get_clientsession(self.hass), host)
 
     @staticmethod
     @callback
-    def async_get_options_flow(
-        config_entry: config_entries.ConfigEntry,
-    ) -> OptionsFlowHandler:
+    def async_get_options_flow(config_entry: ConfigEntry) -> OptionsFlowHandler:
         """Get the options flow for this handler."""
         return OptionsFlowHandler(config_entry)
 
     @classmethod
     @callback
-    def async_supports_options_flow(
-        cls, config_entry: config_entries.ConfigEntry
-    ) -> bool:
+    def async_supports_options_flow(cls, config_entry: ConfigEntry) -> bool:
         """Return options flow support for this handler."""
         return config_entry.data.get("gen") == 2 and not config_entry.data.get(
             CONF_SLEEP_PERIOD
         )
 
 
-class OptionsFlowHandler(config_entries.OptionsFlow):
+class OptionsFlowHandler(OptionsFlow):
     """Handle the option flow for shelly."""
 
-    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
+    def __init__(self, config_entry: ConfigEntry) -> None:
         """Initialize options flow."""
         self.config_entry = config_entry
 
@@ -407,8 +402,8 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
                         default=self.config_entry.options.get(
                             CONF_BLE_SCANNER_MODE, BLEScannerMode.DISABLED
                         ),
-                    ): selector.SelectSelector(
-                        selector.SelectSelectorConfig(options=BLE_SCANNER_OPTIONS),
+                    ): SelectSelector(
+                        SelectSelectorConfig(options=BLE_SCANNER_OPTIONS),
                     ),
                 }
             ),
diff --git a/homeassistant/components/shelly/coordinator.py b/homeassistant/components/shelly/coordinator.py
index 1c4ba7d4763..84964070686 100644
--- a/homeassistant/components/shelly/coordinator.py
+++ b/homeassistant/components/shelly/coordinator.py
@@ -14,12 +14,14 @@ from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCal
 from aioshelly.rpc_device import RpcDevice, UpdateType
 from awesomeversion import AwesomeVersion
 
-from homeassistant import config_entries
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import ATTR_DEVICE_ID, CONF_HOST, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, callback
-from homeassistant.helpers import device_registry
 from homeassistant.helpers.debounce import Debouncer
+from homeassistant.helpers.device_registry import (
+    CONNECTION_NETWORK_MAC,
+    async_get as dr_async_get,
+)
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .bluetooth import async_connect_scanner
@@ -119,11 +121,11 @@ class ShellyCoordinatorBase(DataUpdateCoordinator[None], Generic[_DeviceT]):
 
     def async_setup(self) -> None:
         """Set up the coordinator."""
-        dev_reg = device_registry.async_get(self.hass)
+        dev_reg = dr_async_get(self.hass)
         device_entry = dev_reg.async_get_or_create(
             config_entry_id=self.entry.entry_id,
             name=self.name,
-            connections={(device_registry.CONNECTION_NETWORK_MAC, self.mac)},
+            connections={(CONNECTION_NETWORK_MAC, self.mac)},
             manufacturer="Shelly",
             model=aioshelly.const.MODEL_NAMES.get(self.model, self.model),
             sw_version=self.sw_version,
@@ -563,7 +565,7 @@ def get_block_coordinator_by_device_id(
     hass: HomeAssistant, device_id: str
 ) -> ShellyBlockCoordinator | None:
     """Get a Shelly block device coordinator for the given device id."""
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr_async_get(hass)
     if device := dev_reg.async_get(device_id):
         for config_entry in device.config_entries:
             if not (entry_data := get_entry_data(hass).get(config_entry)):
@@ -579,7 +581,7 @@ def get_rpc_coordinator_by_device_id(
     hass: HomeAssistant, device_id: str
 ) -> ShellyRpcCoordinator | None:
     """Get a Shelly RPC device coordinator for the given device id."""
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr_async_get(hass)
     if device := dev_reg.async_get(device_id):
         for config_entry in device.config_entries:
             if not (entry_data := get_entry_data(hass).get(config_entry)):
@@ -591,14 +593,12 @@ def get_rpc_coordinator_by_device_id(
     return None
 
 
-async def async_reconnect_soon(
-    hass: HomeAssistant, entry: config_entries.ConfigEntry
-) -> None:
+async def async_reconnect_soon(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Try to reconnect soon."""
     if (
         not entry.data.get(CONF_SLEEP_PERIOD)
         and not hass.is_stopping
-        and entry.state == config_entries.ConfigEntryState.LOADED
+        and entry.state == ConfigEntryState.LOADED
         and (entry_data := get_entry_data(hass).get(entry.entry_id))
         and (coordinator := entry_data.rpc)
     ):
diff --git a/homeassistant/components/shelly/entity.py b/homeassistant/components/shelly/entity.py
index 2ab0af8f18e..4811334b285 100644
--- a/homeassistant/components/shelly/entity.py
+++ b/homeassistant/components/shelly/entity.py
@@ -12,10 +12,14 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import device_registry, entity, entity_registry
-from homeassistant.helpers.entity import DeviceInfo, EntityDescription
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.entity_registry import RegistryEntry
+from homeassistant.helpers.entity_registry import (
+    RegistryEntry,
+    async_entries_for_config_entry,
+    async_get as er_async_get,
+)
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -115,10 +119,8 @@ def async_restore_block_attribute_entities(
     """Restore block attributes entities."""
     entities = []
 
-    ent_reg = entity_registry.async_get(hass)
-    entries = entity_registry.async_entries_for_config_entry(
-        ent_reg, config_entry.entry_id
-    )
+    ent_reg = er_async_get(hass)
+    entries = async_entries_for_config_entry(ent_reg, config_entry.entry_id)
 
     domain = sensor_class.__module__.split(".")[-1]
 
@@ -228,10 +230,8 @@ def async_restore_rpc_attribute_entities(
     """Restore block attributes entities."""
     entities = []
 
-    ent_reg = entity_registry.async_get(hass)
-    entries = entity_registry.async_entries_for_config_entry(
-        ent_reg, config_entry.entry_id
-    )
+    ent_reg = er_async_get(hass)
+    entries = async_entries_for_config_entry(ent_reg, config_entry.entry_id)
 
     domain = sensor_class.__module__.split(".")[-1]
 
@@ -321,7 +321,7 @@ class ShellyBlockEntity(CoordinatorEntity[ShellyBlockCoordinator]):
         self._attr_name = get_block_entity_name(coordinator.device, block)
         self._attr_should_poll = False
         self._attr_device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
+            connections={(CONNECTION_NETWORK_MAC, coordinator.mac)}
         )
         self._attr_unique_id = f"{coordinator.mac}-{block.description}"
 
@@ -363,7 +363,7 @@ class ShellyRpcEntity(CoordinatorEntity[ShellyRpcCoordinator]):
         self.key = key
         self._attr_should_poll = False
         self._attr_device_info = {
-            "connections": {(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
+            "connections": {(CONNECTION_NETWORK_MAC, coordinator.mac)}
         }
         self._attr_unique_id = f"{coordinator.mac}-{key}"
         self._attr_name = get_rpc_entity_name(coordinator.device, key)
@@ -412,7 +412,7 @@ class ShellyRpcEntity(CoordinatorEntity[ShellyRpcCoordinator]):
             self.coordinator.entry.async_start_reauth(self.hass)
 
 
-class ShellyBlockAttributeEntity(ShellyBlockEntity, entity.Entity):
+class ShellyBlockAttributeEntity(ShellyBlockEntity, Entity):
     """Helper class to represent a block attribute."""
 
     entity_description: BlockEntityDescription
@@ -482,7 +482,7 @@ class ShellyRestAttributeEntity(CoordinatorEntity[ShellyBlockCoordinator]):
         )
         self._attr_unique_id = f"{coordinator.mac}-{attribute}"
         self._attr_device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
+            connections={(CONNECTION_NETWORK_MAC, coordinator.mac)}
         )
         self._last_value = None
 
@@ -501,7 +501,7 @@ class ShellyRestAttributeEntity(CoordinatorEntity[ShellyBlockCoordinator]):
         return self._last_value
 
 
-class ShellyRpcAttributeEntity(ShellyRpcEntity, entity.Entity):
+class ShellyRpcAttributeEntity(ShellyRpcEntity, Entity):
     """Helper class to represent a rpc attribute."""
 
     entity_description: RpcEntityDescription
@@ -575,7 +575,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
 
         self._attr_should_poll = False
         self._attr_device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
+            connections={(CONNECTION_NETWORK_MAC, coordinator.mac)}
         )
 
         if block is not None:
@@ -658,7 +658,7 @@ class ShellySleepingRpcAttributeEntity(ShellyRpcAttributeEntity, RestoreEntity):
 
         self._attr_should_poll = False
         self._attr_device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
+            connections={(CONNECTION_NETWORK_MAC, coordinator.mac)}
         )
         self._attr_unique_id = (
             self._attr_unique_id
diff --git a/homeassistant/components/shelly/utils.py b/homeassistant/components/shelly/utils.py
index edfa1d284ed..4183b527596 100644
--- a/homeassistant/components/shelly/utils.py
+++ b/homeassistant/components/shelly/utils.py
@@ -13,7 +13,13 @@ from homeassistant.components.http import HomeAssistantView
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry, entity_registry, singleton
+from homeassistant.helpers import singleton
+from homeassistant.helpers.device_registry import (
+    CONNECTION_NETWORK_MAC,
+    async_get as dr_async_get,
+    format_mac,
+)
+from homeassistant.helpers.entity_registry import async_get as er_async_get
 from homeassistant.helpers.typing import EventType
 from homeassistant.util.dt import utcnow
 
@@ -36,7 +42,7 @@ def async_remove_shelly_entity(
     hass: HomeAssistant, domain: str, unique_id: str
 ) -> None:
     """Remove a Shelly entity."""
-    entity_reg = entity_registry.async_get(hass)
+    entity_reg = er_async_get(hass)
     entity_id = entity_reg.async_get_entity_id(domain, DOMAIN, unique_id)
     if entity_id:
         LOGGER.debug("Removing entity: %s", entity_id)
@@ -393,19 +399,12 @@ def device_update_info(
 
     assert entry.unique_id
 
-    dev_registry = device_registry.async_get(hass)
-    if device := dev_registry.async_get_device(
+    dev_reg = dr_async_get(hass)
+    if device := dev_reg.async_get_device(
         identifiers={(DOMAIN, entry.entry_id)},
-        connections={
-            (
-                device_registry.CONNECTION_NETWORK_MAC,
-                device_registry.format_mac(entry.unique_id),
-            )
-        },
+        connections={(CONNECTION_NETWORK_MAC, format_mac(entry.unique_id))},
     ):
-        dev_registry.async_update_device(
-            device.id, sw_version=shellydevice.firmware_version
-        )
+        dev_reg.async_update_device(device.id, sw_version=shellydevice.firmware_version)
 
 
 def brightness_to_percentage(brightness: int) -> int:
diff --git a/tests/components/shelly/__init__.py b/tests/components/shelly/__init__.py
index b58b147a9a3..346eed45baf 100644
--- a/tests/components/shelly/__init__.py
+++ b/tests/components/shelly/__init__.py
@@ -18,7 +18,7 @@ from homeassistant.components.shelly.const import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, format_mac
 from homeassistant.helpers.entity_registry import async_get
 from homeassistant.util import dt
 
@@ -118,10 +118,5 @@ def register_device(device_reg, config_entry: ConfigEntry):
     """Register Shelly device."""
     device_reg.async_get_or_create(
         config_entry_id=config_entry.entry_id,
-        connections={
-            (
-                device_registry.CONNECTION_NETWORK_MAC,
-                device_registry.format_mac(MOCK_MAC),
-            )
-        },
+        connections={(CONNECTION_NETWORK_MAC, format_mac(MOCK_MAC))},
     )
diff --git a/tests/components/shelly/test_config_flow.py b/tests/components/shelly/test_config_flow.py
index fbaf2d98ba2..7a8ffc87e52 100644
--- a/tests/components/shelly/test_config_flow.py
+++ b/tests/components/shelly/test_config_flow.py
@@ -66,7 +66,7 @@ async def test_form(hass, gen):
     assert result["errors"] == {}
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False, "gen": gen},
     ), patch(
         "aioshelly.block_device.BlockDevice.create",
@@ -123,7 +123,7 @@ async def test_title_without_name(hass):
     settings["device"] = settings["device"].copy()
     settings["device"]["hostname"] = "shelly1pm-12345"
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False},
     ), patch(
         "aioshelly.block_device.BlockDevice.create",
@@ -174,7 +174,7 @@ async def test_form_auth(hass, test_data):
     assert result["errors"] == {}
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": True, "gen": gen},
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -237,7 +237,7 @@ async def test_form_errors_get_info(hass, error):
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
-    with patch("aioshelly.common.get_info", side_effect=exc):
+    with patch("homeassistant.components.shelly.config_flow.get_info", side_effect=exc):
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {"host": "1.1.1.1"},
@@ -253,7 +253,7 @@ async def test_form_missing_model_key(hass):
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "auth": False, "gen": "2"},
     ), patch(
         "aioshelly.rpc_device.RpcDevice.create",
@@ -283,7 +283,7 @@ async def test_form_missing_model_key_auth_enabled(hass):
     assert result["errors"] == {}
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "auth": True, "gen": 2},
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -316,7 +316,7 @@ async def test_form_missing_model_key_zeroconf(hass, caplog):
     """Test we handle missing Shelly model key via zeroconf."""
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "auth": False, "gen": 2},
     ), patch(
         "aioshelly.rpc_device.RpcDevice.create",
@@ -355,7 +355,8 @@ async def test_form_errors_test_connection(hass, error):
     )
 
     with patch(
-        "aioshelly.common.get_info", return_value={"mac": "test-mac", "auth": False}
+        "homeassistant.components.shelly.config_flow.get_info",
+        return_value={"mac": "test-mac", "auth": False},
     ), patch(
         "aioshelly.block_device.BlockDevice.create", new=AsyncMock(side_effect=exc)
     ):
@@ -381,7 +382,7 @@ async def test_form_already_configured(hass):
     )
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False},
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -416,7 +417,7 @@ async def test_user_setup_ignored_device(hass):
     settings["fw"] = "20201124-092534/v1.9.0@57ac4ad8"
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False},
     ), patch(
         "aioshelly.block_device.BlockDevice.create",
@@ -452,7 +453,10 @@ async def test_form_firmware_unsupported(hass):
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
-    with patch("aioshelly.common.get_info", side_effect=FirmwareUnsupported):
+    with patch(
+        "homeassistant.components.shelly.config_flow.get_info",
+        side_effect=FirmwareUnsupported,
+    ):
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {"host": "1.1.1.1"},
@@ -478,7 +482,7 @@ async def test_form_auth_errors_test_connection_gen1(hass, error):
     )
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "auth": True},
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -514,7 +518,7 @@ async def test_form_auth_errors_test_connection_gen2(hass, error):
     )
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "auth": True, "gen": 2},
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -543,7 +547,9 @@ async def test_form_auth_errors_test_connection_gen2(hass, error):
 async def test_zeroconf(hass, gen, get_info):
     """Test we get the form."""
 
-    with patch("aioshelly.common.get_info", return_value=get_info), patch(
+    with patch(
+        "homeassistant.components.shelly.config_flow.get_info", return_value=get_info
+    ), patch(
         "aioshelly.block_device.BlockDevice.create",
         new=AsyncMock(return_value=Mock(model="SHSW-1", settings=MOCK_SETTINGS)),
     ), patch(
@@ -598,7 +604,7 @@ async def test_zeroconf_sleeping_device(hass):
     """Test sleeping device configuration via zeroconf."""
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={
             "mac": "test-mac",
             "type": "SHSW-1",
@@ -662,7 +668,7 @@ async def test_zeroconf_sleeping_device(hass):
 async def test_zeroconf_sleeping_device_error(hass):
     """Test sleeping device configuration via zeroconf with error."""
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={
             "mac": "test-mac",
             "type": "SHSW-1",
@@ -691,7 +697,7 @@ async def test_zeroconf_already_configured(hass):
     entry.add_to_hass(hass)
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False},
     ):
         result = await hass.config_entries.flow.async_init(
@@ -718,7 +724,7 @@ async def test_zeroconf_ignored(hass):
     entry.add_to_hass(hass)
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False},
     ):
         result = await hass.config_entries.flow.async_init(
@@ -739,7 +745,7 @@ async def test_zeroconf_with_wifi_ap_ip(hass):
     entry.add_to_hass(hass)
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False},
     ):
         result = await hass.config_entries.flow.async_init(
@@ -756,7 +762,10 @@ async def test_zeroconf_with_wifi_ap_ip(hass):
 
 async def test_zeroconf_firmware_unsupported(hass):
     """Test we abort if device firmware is unsupported."""
-    with patch("aioshelly.common.get_info", side_effect=FirmwareUnsupported):
+    with patch(
+        "homeassistant.components.shelly.config_flow.get_info",
+        side_effect=FirmwareUnsupported,
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN,
             data=DISCOVERY_INFO,
@@ -769,7 +778,10 @@ async def test_zeroconf_firmware_unsupported(hass):
 
 async def test_zeroconf_cannot_connect(hass):
     """Test we get the form."""
-    with patch("aioshelly.common.get_info", side_effect=DeviceConnectionError):
+    with patch(
+        "homeassistant.components.shelly.config_flow.get_info",
+        side_effect=DeviceConnectionError,
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN,
             data=DISCOVERY_INFO,
@@ -783,7 +795,7 @@ async def test_zeroconf_require_auth(hass):
     """Test zeroconf if auth is required."""
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": True},
     ):
         result = await hass.config_entries.flow.async_init(
@@ -844,7 +856,7 @@ async def test_reauth_successful(hass, test_data):
     entry.add_to_hass(hass)
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": True, "gen": gen},
     ), patch(
         "aioshelly.block_device.BlockDevice.create",
@@ -898,7 +910,7 @@ async def test_reauth_unsuccessful(hass, test_data):
     entry.add_to_hass(hass)
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "test-mac", "type": "SHSW-1", "auth": True, "gen": gen},
     ), patch(
         "aioshelly.block_device.BlockDevice.create",
@@ -937,7 +949,7 @@ async def test_reauth_get_info_error(hass, error):
     entry.add_to_hass(hass)
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         side_effect=error,
     ):
         result = await hass.config_entries.flow.async_init(
@@ -1140,7 +1152,7 @@ async def test_zeroconf_already_configured_triggers_refresh_mac_in_name(
     assert len(mock_rpc_device.initialize.mock_calls) == 1
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "", "type": "SHSW-1", "auth": False},
     ):
         result = await hass.config_entries.flow.async_init(
@@ -1172,7 +1184,7 @@ async def test_zeroconf_already_configured_triggers_refresh(
     assert len(mock_rpc_device.initialize.mock_calls) == 1
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "AABBCCDDEEFF", "type": "SHSW-1", "auth": False},
     ):
         result = await hass.config_entries.flow.async_init(
@@ -1207,7 +1219,7 @@ async def test_zeroconf_sleeping_device_not_triggers_refresh(
     assert "online, resuming setup" in caplog.text
 
     with patch(
-        "aioshelly.common.get_info",
+        "homeassistant.components.shelly.config_flow.get_info",
         return_value={"mac": "AABBCCDDEEFF", "type": "SHSW-1", "auth": False},
     ):
         result = await hass.config_entries.flow.async_init(
diff --git a/tests/components/shelly/test_device_trigger.py b/tests/components/shelly/test_device_trigger.py
index 13ccf3f843e..31e54545b28 100644
--- a/tests/components/shelly/test_device_trigger.py
+++ b/tests/components/shelly/test_device_trigger.py
@@ -14,8 +14,8 @@ from homeassistant.components.shelly.const import (
     EVENT_SHELLY_CLICK,
 )
 from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_PLATFORM, CONF_TYPE
-from homeassistant.helpers import device_registry
 from homeassistant.helpers.device_registry import (
+    CONNECTION_NETWORK_MAC,
     async_entries_for_config_entry,
     async_get as async_get_dev_reg,
 )
@@ -151,7 +151,7 @@ async def test_get_triggers_for_invalid_device_id(hass, device_reg, mock_block_d
     config_entry.add_to_hass(hass)
     invalid_device = device_reg.async_get_or_create(
         config_entry_id=config_entry.entry_id,
-        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        connections={(CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
     )
 
     with pytest.raises(InvalidDeviceAutomationConfig):
diff --git a/tests/components/shelly/test_init.py b/tests/components/shelly/test_init.py
index 3675186b9ba..707e8d5cfb1 100644
--- a/tests/components/shelly/test_init.py
+++ b/tests/components/shelly/test_init.py
@@ -13,7 +13,7 @@ from homeassistant.components.shelly.const import (
 )
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.const import STATE_ON, STATE_UNAVAILABLE
-from homeassistant.helpers import device_registry
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, format_mac
 from homeassistant.setup import async_setup_component
 
 from . import MOCK_MAC, init_integration
@@ -43,12 +43,7 @@ async def test_shared_device_mac(
     config_entry.add_to_hass(hass)
     device_reg.async_get_or_create(
         config_entry_id=config_entry.entry_id,
-        connections={
-            (
-                device_registry.CONNECTION_NETWORK_MAC,
-                device_registry.format_mac(MOCK_MAC),
-            )
-        },
+        connections={(CONNECTION_NETWORK_MAC, format_mac(MOCK_MAC))},
     )
     await init_integration(hass, gen, sleep_period=1000)
     assert "will resume when device is online" in caplog.text
@@ -117,12 +112,7 @@ async def test_sleeping_block_device_online(
     config_entry.add_to_hass(hass)
     device_reg.async_get_or_create(
         config_entry_id=config_entry.entry_id,
-        connections={
-            (
-                device_registry.CONNECTION_NETWORK_MAC,
-                device_registry.format_mac(MOCK_MAC),
-            )
-        },
+        connections={(CONNECTION_NETWORK_MAC, format_mac(MOCK_MAC))},
     )
 
     entry = await init_integration(hass, 1, sleep_period=entry_sleep)