diff --git a/homeassistant/components/aosmith/__init__.py b/homeassistant/components/aosmith/__init__.py index cac746e189e..b75a4ad7295 100644 --- a/homeassistant/components/aosmith/__init__.py +++ b/homeassistant/components/aosmith/__init__.py @@ -8,10 +8,10 @@ from py_aosmith import AOSmithAPIClient from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform from homeassistant.core import HomeAssistant -from homeassistant.helpers import aiohttp_client +from homeassistant.helpers import aiohttp_client, device_registry as dr from .const import DOMAIN -from .coordinator import AOSmithCoordinator +from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.WATER_HEATER] @@ -20,8 +20,9 @@ PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.WATER_HEATER] class AOSmithData: """Data for the A. O. Smith integration.""" - coordinator: AOSmithCoordinator client: AOSmithAPIClient + status_coordinator: AOSmithStatusCoordinator + energy_coordinator: AOSmithEnergyCoordinator async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: @@ -31,13 +32,32 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: session = aiohttp_client.async_get_clientsession(hass) client = AOSmithAPIClient(email, password, session) - coordinator = AOSmithCoordinator(hass, client) - # Fetch initial data so we have data when entities subscribe - await coordinator.async_config_entry_first_refresh() + status_coordinator = AOSmithStatusCoordinator(hass, client) + await status_coordinator.async_config_entry_first_refresh() + + device_registry = dr.async_get(hass) + for junction_id, status_data in status_coordinator.data.items(): + device_registry.async_get_or_create( + config_entry_id=entry.entry_id, + identifiers={(DOMAIN, junction_id)}, + manufacturer="A. O. Smith", + name=status_data.get("name"), + model=status_data.get("model"), + serial_number=status_data.get("serial"), + suggested_area=status_data.get("install", {}).get("location"), + sw_version=status_data.get("data", {}).get("firmwareVersion"), + ) + + energy_coordinator = AOSmithEnergyCoordinator( + hass, client, list(status_coordinator.data) + ) + await energy_coordinator.async_config_entry_first_refresh() hass.data.setdefault(DOMAIN, {})[entry.entry_id] = AOSmithData( - coordinator=coordinator, client=client + client, + status_coordinator, + energy_coordinator, ) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) diff --git a/homeassistant/components/aosmith/const.py b/homeassistant/components/aosmith/const.py index e79b993182e..c0c693e0dac 100644 --- a/homeassistant/components/aosmith/const.py +++ b/homeassistant/components/aosmith/const.py @@ -15,6 +15,9 @@ REGULAR_INTERVAL = timedelta(seconds=30) # Update interval to be used while a mode or setpoint change is in progress. FAST_INTERVAL = timedelta(seconds=1) +# Update interval to be used for energy usage data. +ENERGY_USAGE_INTERVAL = timedelta(minutes=10) + HOT_WATER_STATUS_MAP = { "LOW": "low", "MEDIUM": "medium", diff --git a/homeassistant/components/aosmith/coordinator.py b/homeassistant/components/aosmith/coordinator.py index bdd144569dd..7d6053cc86e 100644 --- a/homeassistant/components/aosmith/coordinator.py +++ b/homeassistant/components/aosmith/coordinator.py @@ -1,5 +1,4 @@ """The data update coordinator for the A. O. Smith integration.""" - import logging from typing import Any @@ -13,13 +12,13 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed -from .const import DOMAIN, FAST_INTERVAL, REGULAR_INTERVAL +from .const import DOMAIN, ENERGY_USAGE_INTERVAL, FAST_INTERVAL, REGULAR_INTERVAL _LOGGER = logging.getLogger(__name__) -class AOSmithCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]): - """Custom data update coordinator for A. O. Smith integration.""" +class AOSmithStatusCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]): + """Coordinator for device status, updating with a frequent interval.""" def __init__(self, hass: HomeAssistant, client: AOSmithAPIClient) -> None: """Initialize the coordinator.""" @@ -27,7 +26,7 @@ class AOSmithCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]): self.client = client async def _async_update_data(self) -> dict[str, dict[str, Any]]: - """Fetch latest data from API.""" + """Fetch latest data from the device status endpoint.""" try: devices = await self.client.get_devices() except AOSmithInvalidCredentialsException as err: @@ -49,3 +48,36 @@ class AOSmithCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]): self.update_interval = REGULAR_INTERVAL return {device.get("junctionId"): device for device in devices} + + +class AOSmithEnergyCoordinator(DataUpdateCoordinator[dict[str, float]]): + """Coordinator for energy usage data, updating with a slower interval.""" + + def __init__( + self, + hass: HomeAssistant, + client: AOSmithAPIClient, + junction_ids: list[str], + ) -> None: + """Initialize the coordinator.""" + super().__init__( + hass, _LOGGER, name=DOMAIN, update_interval=ENERGY_USAGE_INTERVAL + ) + self.client = client + self.junction_ids = junction_ids + + async def _async_update_data(self) -> dict[str, float]: + """Fetch latest data from the energy usage endpoint.""" + energy_usage_by_junction_id: dict[str, float] = {} + + for junction_id in self.junction_ids: + try: + energy_usage = await self.client.get_energy_use_data(junction_id) + except AOSmithInvalidCredentialsException as err: + raise ConfigEntryAuthFailed from err + except AOSmithUnknownException as err: + raise UpdateFailed(f"Error communicating with API: {err}") from err + + energy_usage_by_junction_id[junction_id] = energy_usage.get("lifetimeKwh") + + return energy_usage_by_junction_id diff --git a/homeassistant/components/aosmith/entity.py b/homeassistant/components/aosmith/entity.py index 20061ca36b9..107e5d7e944 100644 --- a/homeassistant/components/aosmith/entity.py +++ b/homeassistant/components/aosmith/entity.py @@ -1,5 +1,5 @@ """The base entity for the A. O. Smith integration.""" - +from typing import TypeVar from py_aosmith import AOSmithAPIClient @@ -7,28 +7,35 @@ from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN -from .coordinator import AOSmithCoordinator +from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator + +_AOSmithCoordinatorT = TypeVar( + "_AOSmithCoordinatorT", bound=AOSmithStatusCoordinator | AOSmithEnergyCoordinator +) -class AOSmithEntity(CoordinatorEntity[AOSmithCoordinator]): +class AOSmithEntity(CoordinatorEntity[_AOSmithCoordinatorT]): """Base entity for A. O. Smith.""" _attr_has_entity_name = True - def __init__(self, coordinator: AOSmithCoordinator, junction_id: str) -> None: + def __init__(self, coordinator: _AOSmithCoordinatorT, junction_id: str) -> None: """Initialize the entity.""" super().__init__(coordinator) self.junction_id = junction_id self._attr_device_info = DeviceInfo( - manufacturer="A. O. Smith", - name=self.device.get("name"), - model=self.device.get("model"), - serial_number=self.device.get("serial"), - suggested_area=self.device.get("install", {}).get("location"), identifiers={(DOMAIN, junction_id)}, - sw_version=self.device.get("data", {}).get("firmwareVersion"), ) + @property + def client(self) -> AOSmithAPIClient: + """Shortcut to get the API client.""" + return self.coordinator.client + + +class AOSmithStatusEntity(AOSmithEntity[AOSmithStatusCoordinator]): + """Base entity for entities that use data from the status coordinator.""" + @property def device(self): """Shortcut to get the device status from the coordinator data.""" @@ -40,12 +47,16 @@ class AOSmithEntity(CoordinatorEntity[AOSmithCoordinator]): device = self.device return None if device is None else device.get("data", {}) - @property - def client(self) -> AOSmithAPIClient: - """Shortcut to get the API client.""" - return self.coordinator.client - @property def available(self) -> bool: """Return True if entity is available.""" return super().available and self.device_data.get("isOnline") is True + + +class AOSmithEnergyEntity(AOSmithEntity[AOSmithEnergyCoordinator]): + """Base entity for entities that use data from the energy coordinator.""" + + @property + def energy_usage(self) -> float | None: + """Shortcut to get the energy usage from the coordinator data.""" + return self.coordinator.data.get(self.junction_id) diff --git a/homeassistant/components/aosmith/manifest.json b/homeassistant/components/aosmith/manifest.json index 2e3a459d7e2..895b03cf7fd 100644 --- a/homeassistant/components/aosmith/manifest.json +++ b/homeassistant/components/aosmith/manifest.json @@ -5,6 +5,5 @@ "config_flow": true, "documentation": "https://www.home-assistant.io/integrations/aosmith", "iot_class": "cloud_polling", - "quality_scale": "platinum", "requirements": ["py-aosmith==1.0.1"] } diff --git a/homeassistant/components/aosmith/sensor.py b/homeassistant/components/aosmith/sensor.py index 78c6f32232a..b0606d2dca4 100644 --- a/homeassistant/components/aosmith/sensor.py +++ b/homeassistant/components/aosmith/sensor.py @@ -8,26 +8,28 @@ from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, SensorEntityDescription, + SensorStateClass, ) from homeassistant.config_entries import ConfigEntry +from homeassistant.const import UnitOfEnergy from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import AOSmithData from .const import DOMAIN, HOT_WATER_STATUS_MAP -from .coordinator import AOSmithCoordinator -from .entity import AOSmithEntity +from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator +from .entity import AOSmithEnergyEntity, AOSmithStatusEntity @dataclass(frozen=True, kw_only=True) -class AOSmithSensorEntityDescription(SensorEntityDescription): - """Define sensor entity description class.""" +class AOSmithStatusSensorEntityDescription(SensorEntityDescription): + """Entity description class for sensors using data from the status coordinator.""" value_fn: Callable[[dict[str, Any]], str | int | None] -ENTITY_DESCRIPTIONS: tuple[AOSmithSensorEntityDescription, ...] = ( - AOSmithSensorEntityDescription( +STATUS_ENTITY_DESCRIPTIONS: tuple[AOSmithStatusSensorEntityDescription, ...] = ( + AOSmithStatusSensorEntityDescription( key="hot_water_availability", translation_key="hot_water_availability", icon="mdi:water-thermometer", @@ -47,21 +49,26 @@ async def async_setup_entry( data: AOSmithData = hass.data[DOMAIN][entry.entry_id] async_add_entities( - AOSmithSensorEntity(data.coordinator, description, junction_id) - for description in ENTITY_DESCRIPTIONS - for junction_id in data.coordinator.data + AOSmithStatusSensorEntity(data.status_coordinator, description, junction_id) + for description in STATUS_ENTITY_DESCRIPTIONS + for junction_id in data.status_coordinator.data + ) + + async_add_entities( + AOSmithEnergySensorEntity(data.energy_coordinator, junction_id) + for junction_id in data.status_coordinator.data ) -class AOSmithSensorEntity(AOSmithEntity, SensorEntity): - """The sensor entity for the A. O. Smith integration.""" +class AOSmithStatusSensorEntity(AOSmithStatusEntity, SensorEntity): + """Class for sensor entities that use data from the status coordinator.""" - entity_description: AOSmithSensorEntityDescription + entity_description: AOSmithStatusSensorEntityDescription def __init__( self, - coordinator: AOSmithCoordinator, - description: AOSmithSensorEntityDescription, + coordinator: AOSmithStatusCoordinator, + description: AOSmithStatusSensorEntityDescription, junction_id: str, ) -> None: """Initialize the entity.""" @@ -73,3 +80,27 @@ class AOSmithSensorEntity(AOSmithEntity, SensorEntity): def native_value(self) -> str | int | None: """Return the state of the sensor.""" return self.entity_description.value_fn(self.device) + + +class AOSmithEnergySensorEntity(AOSmithEnergyEntity, SensorEntity): + """Class for the energy sensor entity.""" + + _attr_translation_key = "energy_usage" + _attr_device_class = SensorDeviceClass.ENERGY + _attr_state_class = SensorStateClass.TOTAL_INCREASING + _attr_native_unit_of_measurement = UnitOfEnergy.KILO_WATT_HOUR + _attr_suggested_display_precision = 1 + + def __init__( + self, + coordinator: AOSmithEnergyCoordinator, + junction_id: str, + ) -> None: + """Initialize the entity.""" + super().__init__(coordinator, junction_id) + self._attr_unique_id = f"energy_usage_{junction_id}" + + @property + def native_value(self) -> float | None: + """Return the state of the sensor.""" + return self.energy_usage diff --git a/homeassistant/components/aosmith/strings.json b/homeassistant/components/aosmith/strings.json index 0f1fcfc1744..0ca4e2e9094 100644 --- a/homeassistant/components/aosmith/strings.json +++ b/homeassistant/components/aosmith/strings.json @@ -34,6 +34,9 @@ "medium": "Medium", "high": "High" } + }, + "energy_usage": { + "name": "Energy usage" } } } diff --git a/homeassistant/components/aosmith/water_heater.py b/homeassistant/components/aosmith/water_heater.py index 8002373573f..8c42048d439 100644 --- a/homeassistant/components/aosmith/water_heater.py +++ b/homeassistant/components/aosmith/water_heater.py @@ -23,8 +23,8 @@ from .const import ( AOSMITH_MODE_VACATION, DOMAIN, ) -from .coordinator import AOSmithCoordinator -from .entity import AOSmithEntity +from .coordinator import AOSmithStatusCoordinator +from .entity import AOSmithStatusEntity MODE_HA_TO_AOSMITH = { STATE_OFF: AOSMITH_MODE_VACATION, @@ -54,22 +54,24 @@ async def async_setup_entry( """Set up A. O. Smith water heater platform.""" data: AOSmithData = hass.data[DOMAIN][entry.entry_id] - entities = [] - - for junction_id in data.coordinator.data: - entities.append(AOSmithWaterHeaterEntity(data.coordinator, junction_id)) - - async_add_entities(entities) + async_add_entities( + AOSmithWaterHeaterEntity(data.status_coordinator, junction_id) + for junction_id in data.status_coordinator.data + ) -class AOSmithWaterHeaterEntity(AOSmithEntity, WaterHeaterEntity): +class AOSmithWaterHeaterEntity(AOSmithStatusEntity, WaterHeaterEntity): """The water heater entity for the A. O. Smith integration.""" _attr_name = None _attr_temperature_unit = UnitOfTemperature.FAHRENHEIT _attr_min_temp = 95 - def __init__(self, coordinator: AOSmithCoordinator, junction_id: str) -> None: + def __init__( + self, + coordinator: AOSmithStatusCoordinator, + junction_id: str, + ) -> None: """Initialize the entity.""" super().__init__(coordinator, junction_id) self._attr_unique_id = junction_id diff --git a/tests/components/aosmith/conftest.py b/tests/components/aosmith/conftest.py index f0ece65d56f..61c1fc9a562 100644 --- a/tests/components/aosmith/conftest.py +++ b/tests/components/aosmith/conftest.py @@ -10,7 +10,11 @@ from homeassistant.const import CONF_EMAIL, CONF_PASSWORD from homeassistant.core import HomeAssistant from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM -from tests.common import MockConfigEntry, load_json_array_fixture +from tests.common import ( + MockConfigEntry, + load_json_array_fixture, + load_json_object_fixture, +) FIXTURE_USER_INPUT = { CONF_EMAIL: "testemail@example.com", @@ -47,9 +51,13 @@ def get_devices_fixture() -> str: async def mock_client(get_devices_fixture: str) -> Generator[MagicMock, None, None]: """Return a mocked client.""" get_devices_fixture = load_json_array_fixture(f"{get_devices_fixture}.json", DOMAIN) + get_energy_use_fixture = load_json_object_fixture( + "get_energy_use_data.json", DOMAIN + ) client_mock = MagicMock(AOSmithAPIClient) client_mock.get_devices = AsyncMock(return_value=get_devices_fixture) + client_mock.get_energy_use_data = AsyncMock(return_value=get_energy_use_fixture) return client_mock diff --git a/tests/components/aosmith/fixtures/get_energy_use_data.json b/tests/components/aosmith/fixtures/get_energy_use_data.json new file mode 100644 index 00000000000..989ddab5399 --- /dev/null +++ b/tests/components/aosmith/fixtures/get_energy_use_data.json @@ -0,0 +1,19 @@ +{ + "average": 2.7552000000000003, + "graphData": [ + { + "date": "2023-10-30T04:00:00.000Z", + "kwh": 2.01 + }, + { + "date": "2023-10-31T04:00:00.000Z", + "kwh": 1.542 + }, + { + "date": "2023-11-01T04:00:00.000Z", + "kwh": 1.908 + } + ], + "lifetimeKwh": 132.825, + "startDate": "Oct 30" +} diff --git a/tests/components/aosmith/snapshots/test_sensor.ambr b/tests/components/aosmith/snapshots/test_sensor.ambr index 8499a98c8e5..d4376c64a01 100644 --- a/tests/components/aosmith/snapshots/test_sensor.ambr +++ b/tests/components/aosmith/snapshots/test_sensor.ambr @@ -1,5 +1,20 @@ # serializer version: 1 -# name: test_state +# name: test_state[sensor.my_water_heater_energy_usage] + StateSnapshot({ + 'attributes': ReadOnlyDict({ + 'device_class': 'energy', + 'friendly_name': 'My water heater Energy usage', + 'state_class': , + 'unit_of_measurement': , + }), + 'context': , + 'entity_id': 'sensor.my_water_heater_energy_usage', + 'last_changed': , + 'last_updated': , + 'state': '132.825', + }) +# --- +# name: test_state[sensor.my_water_heater_hot_water_availability] StateSnapshot({ 'attributes': ReadOnlyDict({ 'device_class': 'enum', diff --git a/tests/components/aosmith/test_config_flow.py b/tests/components/aosmith/test_config_flow.py index 5d3e986e05e..d6cf1655b14 100644 --- a/tests/components/aosmith/test_config_flow.py +++ b/tests/components/aosmith/test_config_flow.py @@ -1,4 +1,5 @@ """Test the A. O. Smith config flow.""" +from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch from freezegun.api import FrozenDateTimeFactory @@ -6,7 +7,11 @@ from py_aosmith import AOSmithInvalidCredentialsException import pytest from homeassistant import config_entries -from homeassistant.components.aosmith.const import DOMAIN, REGULAR_INTERVAL +from homeassistant.components.aosmith.const import ( + DOMAIN, + ENERGY_USAGE_INTERVAL, + REGULAR_INTERVAL, +) from homeassistant.config_entries import ConfigEntryState from homeassistant.const import CONF_EMAIL, CONF_PASSWORD from homeassistant.core import HomeAssistant @@ -87,21 +92,30 @@ async def test_form_exception( assert len(mock_setup_entry.mock_calls) == 1 +@pytest.mark.parametrize( + ("api_method", "wait_interval"), + [ + ("get_devices", REGULAR_INTERVAL), + ("get_energy_use_data", ENERGY_USAGE_INTERVAL), + ], +) async def test_reauth_flow( freezer: FrozenDateTimeFactory, hass: HomeAssistant, init_integration: MockConfigEntry, mock_client: MagicMock, + api_method: str, + wait_interval: timedelta, ) -> None: """Test reauth works.""" entries = hass.config_entries.async_entries(DOMAIN) assert len(entries) == 1 assert entries[0].state is ConfigEntryState.LOADED - mock_client.get_devices.side_effect = AOSmithInvalidCredentialsException( + getattr(mock_client, api_method).side_effect = AOSmithInvalidCredentialsException( "Authentication error" ) - freezer.tick(REGULAR_INTERVAL) + freezer.tick(wait_interval) async_fire_time_changed(hass) await hass.async_block_till_done() @@ -112,6 +126,9 @@ async def test_reauth_flow( with patch( "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices", return_value=[], + ), patch( + "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_energy_use_data", + return_value=[], ), patch("homeassistant.components.aosmith.async_setup_entry", return_value=True): result2 = await hass.config_entries.flow.async_configure( flows[0]["flow_id"], diff --git a/tests/components/aosmith/test_init.py b/tests/components/aosmith/test_init.py index 8ab699e6f1c..463932e930a 100644 --- a/tests/components/aosmith/test_init.py +++ b/tests/components/aosmith/test_init.py @@ -15,7 +15,11 @@ from homeassistant.components.aosmith.const import ( from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant -from tests.common import MockConfigEntry, async_fire_time_changed +from tests.common import ( + MockConfigEntry, + async_fire_time_changed, + load_json_array_fixture, +) async def test_config_entry_setup(init_integration: MockConfigEntry) -> None: @@ -25,10 +29,10 @@ async def test_config_entry_setup(init_integration: MockConfigEntry) -> None: assert mock_config_entry.state is ConfigEntryState.LOADED -async def test_config_entry_not_ready( +async def test_config_entry_not_ready_get_devices_error( hass: HomeAssistant, mock_config_entry: MockConfigEntry ) -> None: - """Test the config entry not ready.""" + """Test the config entry not ready when get_devices fails.""" mock_config_entry.add_to_hass(hass) with patch( @@ -41,6 +45,28 @@ async def test_config_entry_not_ready( assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY +async def test_config_entry_not_ready_get_energy_use_data_error( + hass: HomeAssistant, + mock_config_entry: MockConfigEntry, +) -> None: + """Test the config entry not ready when get_energy_use_data fails.""" + mock_config_entry.add_to_hass(hass) + + get_devices_fixture = load_json_array_fixture("get_devices.json", DOMAIN) + + with patch( + "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices", + return_value=get_devices_fixture, + ), patch( + "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_energy_use_data", + side_effect=AOSmithUnknownException("Unknown error"), + ): + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() + + assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY + + @pytest.mark.parametrize( ("get_devices_fixture", "time_to_wait", "expected_call_count"), [ diff --git a/tests/components/aosmith/test_sensor.py b/tests/components/aosmith/test_sensor.py index 99626b09e83..f94dfdb710c 100644 --- a/tests/components/aosmith/test_sensor.py +++ b/tests/components/aosmith/test_sensor.py @@ -1,5 +1,6 @@ """Tests for the sensor platform of the A. O. Smith integration.""" +import pytest from syrupy.assertion import SnapshotAssertion from homeassistant.core import HomeAssistant @@ -8,20 +9,42 @@ from homeassistant.helpers import entity_registry as er from tests.common import MockConfigEntry +@pytest.mark.parametrize( + ("entity_id", "unique_id"), + [ + ( + "sensor.my_water_heater_hot_water_availability", + "hot_water_availability_junctionId", + ), + ("sensor.my_water_heater_energy_usage", "energy_usage_junctionId"), + ], +) async def test_setup( hass: HomeAssistant, entity_registry: er.EntityRegistry, init_integration: MockConfigEntry, + entity_id: str, + unique_id: str, ) -> None: - """Test the setup of the sensor entity.""" - entry = entity_registry.async_get("sensor.my_water_heater_hot_water_availability") + """Test the setup of the sensor entities.""" + entry = entity_registry.async_get(entity_id) assert entry - assert entry.unique_id == "hot_water_availability_junctionId" + assert entry.unique_id == unique_id +@pytest.mark.parametrize( + ("entity_id"), + [ + "sensor.my_water_heater_hot_water_availability", + "sensor.my_water_heater_energy_usage", + ], +) async def test_state( - hass: HomeAssistant, init_integration: MockConfigEntry, snapshot: SnapshotAssertion + hass: HomeAssistant, + init_integration: MockConfigEntry, + snapshot: SnapshotAssertion, + entity_id: str, ) -> None: - """Test the state of the sensor entity.""" - state = hass.states.get("sensor.my_water_heater_hot_water_availability") + """Test the state of the sensor entities.""" + state = hass.states.get(entity_id) assert state == snapshot