Add energy usage sensor to A. O. Smith integration (#105616)

* Add energy usage sensor to A. O. Smith integration

* Address review comments

* Address review comment

* Address review comment

* Create device outside of the entity class

* Address review comment

* remove platinum

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
+    status_coordinator = AOSmithStatusCoordinator(hass, client)
-    await coordinator.async_config_entry_first_refresh()
+    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)

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",

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]]]):
+class AOSmithStatusCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]):
-    """Custom data update coordinator for A. O. Smith integration."""
+    """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

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)

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"]
 }

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 .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator
-from .entity import AOSmithEntity
+from .entity import AOSmithEnergyEntity, AOSmithStatusEntity
 
 
 @dataclass(frozen=True, kw_only=True)
-class AOSmithSensorEntityDescription(SensorEntityDescription):
+class AOSmithStatusSensorEntityDescription(SensorEntityDescription):
-    """Define sensor entity description class."""
+    """Entity description class for sensors using data from the status coordinator."""
 
     value_fn: Callable[[dict[str, Any]], str | int | None]
 
 
-ENTITY_DESCRIPTIONS: tuple[AOSmithSensorEntityDescription, ...] = (
+STATUS_ENTITY_DESCRIPTIONS: tuple[AOSmithStatusSensorEntityDescription, ...] = (
-    AOSmithSensorEntityDescription(
+    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)
+        AOSmithStatusSensorEntity(data.status_coordinator, description, junction_id)
-        for description in ENTITY_DESCRIPTIONS
+        for description in STATUS_ENTITY_DESCRIPTIONS
-        for junction_id in data.coordinator.data
+        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):
+class AOSmithStatusSensorEntity(AOSmithStatusEntity, SensorEntity):
-    """The sensor entity for the A. O. Smith integration."""
+    """Class for sensor entities that use data from the status coordinator."""
 
-    entity_description: AOSmithSensorEntityDescription
+    entity_description: AOSmithStatusSensorEntityDescription
 
     def __init__(
         self,
-        coordinator: AOSmithCoordinator,
+        coordinator: AOSmithStatusCoordinator,
-        description: AOSmithSensorEntityDescription,
+        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

homeassistant/components/aosmith/strings.json

@@ -34,6 +34,9 @@
           "medium": "Medium",
           "high": "High"
         }
+      },
+      "energy_usage": {
+        "name": "Energy usage"
       }
     }
   }

homeassistant/components/aosmith/water_heater.py

@@ -23,8 +23,8 @@ from .const import (
     AOSMITH_MODE_VACATION,
     DOMAIN,
 )
-from .coordinator import AOSmithCoordinator
+from .coordinator import AOSmithStatusCoordinator
-from .entity import AOSmithEntity
+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 = []
+    async_add_entities(
-
+        AOSmithWaterHeaterEntity(data.status_coordinator, junction_id)
-    for junction_id in data.coordinator.data:
+        for junction_id in data.status_coordinator.data
-        entities.append(AOSmithWaterHeaterEntity(data.coordinator, junction_id))
+    )
-
-    async_add_entities(entities)
 
 
-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

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
 

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"
+}

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': <SensorStateClass.TOTAL_INCREASING: 'total_increasing'>,
+      'unit_of_measurement': <UnitOfEnergy.KILO_WATT_HOUR: 'kWh'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.my_water_heater_energy_usage',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '132.825',
+  })
+# ---
+# name: test_state[sensor.my_water_heater_hot_water_availability]
   StateSnapshot({
     'attributes': ReadOnlyDict({
       'device_class': 'enum',

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"],

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"),
     [

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."""
+    """Test the setup of the sensor entities."""
-    entry = entity_registry.async_get("sensor.my_water_heater_hot_water_availability")
+    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."""
+    """Test the state of the sensor entities."""
-    state = hass.states.get("sensor.my_water_heater_hot_water_availability")
+    state = hass.states.get(entity_id)
     assert state == snapshot