diff --git a/tests/components/dsmr/test_mbus_migration.py b/tests/components/dsmr/test_mbus_migration.py index 95def2f66cf..429128c48bb 100644 --- a/tests/components/dsmr/test_mbus_migration.py +++ b/tests/components/dsmr/test_mbus_migration.py @@ -6,13 +6,16 @@ from decimal import Decimal from homeassistant.components.dsmr.const import DOMAIN from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.core import HomeAssistant -from homeassistant.helpers import entity_registry as er +from homeassistant.helpers import device_registry as dr, entity_registry as er from tests.common import MockConfigEntry async def test_migrate_gas_to_mbus( - hass: HomeAssistant, entity_registry: er.EntityRegistry, dsmr_connection_fixture + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + device_registry: dr.DeviceRegistry, + dsmr_connection_fixture, ) -> None: """Test migration of unique_id.""" (connection_factory, transport, protocol) = dsmr_connection_fixture @@ -42,7 +45,6 @@ async def test_migrate_gas_to_mbus( old_unique_id = "37464C4F32313139303333373331_belgium_5min_gas_meter_reading" - device_registry = hass.helpers.device_registry.async_get(hass) device = device_registry.async_get_or_create( config_entry_id=mock_entry.entry_id, identifiers={(DOMAIN, mock_entry.entry_id)}, @@ -108,7 +110,10 @@ async def test_migrate_gas_to_mbus( async def test_migrate_gas_to_mbus_exists( - hass: HomeAssistant, entity_registry: er.EntityRegistry, dsmr_connection_fixture + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + device_registry: dr.DeviceRegistry, + dsmr_connection_fixture, ) -> None: """Test migration of unique_id.""" (connection_factory, transport, protocol) = dsmr_connection_fixture @@ -138,7 +143,6 @@ async def test_migrate_gas_to_mbus_exists( old_unique_id = "37464C4F32313139303333373331_belgium_5min_gas_meter_reading" - device_registry = hass.helpers.device_registry.async_get(hass) device = device_registry.async_get_or_create( config_entry_id=mock_entry.entry_id, identifiers={(DOMAIN, mock_entry.entry_id)},