Add more sensors to PEGELONLINE (#97295)

* add further sensors

* adjust and improve tests

* add device classes were applicable

* fix doc string

* name for ph comes from device class

* use icon from device class for ph sensor

homeassistant/components/pegel_online/coordinator.py

@@ -1,18 +1,17 @@
 """DataUpdateCoordinator for pegel_online."""
 import logging
 
-from aiopegelonline import CONNECT_ERRORS, PegelOnline, Station
+from aiopegelonline import CONNECT_ERRORS, PegelOnline, Station, StationMeasurements
 
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import MIN_TIME_BETWEEN_UPDATES
-from .model import PegelOnlineData
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class PegelOnlineDataUpdateCoordinator(DataUpdateCoordinator[PegelOnlineData]):
+class PegelOnlineDataUpdateCoordinator(DataUpdateCoordinator[StationMeasurements]):
     """DataUpdateCoordinator for the pegel_online integration."""
 
     def __init__(
@@ -28,13 +27,9 @@ class PegelOnlineDataUpdateCoordinator(DataUpdateCoordinator[PegelOnlineData]):
             update_interval=MIN_TIME_BETWEEN_UPDATES,
         )
 
-    async def _async_update_data(self) -> PegelOnlineData:
+    async def _async_update_data(self) -> StationMeasurements:
         """Fetch data from API endpoint."""
         try:
-            water_level = await self.api.async_get_station_measurement(
+            return await self.api.async_get_station_measurements(self.station.uuid)
-                self.station.uuid
-            )
         except CONNECT_ERRORS as err:
             raise UpdateFailed(f"Failed to communicate with API: {err}") from err
-
-        return {"water_level": water_level}

homeassistant/components/pegel_online/manifest.json

@@ -7,5 +7,5 @@
   "integration_type": "service",
   "iot_class": "cloud_polling",
   "loggers": ["aiopegelonline"],
-  "requirements": ["aiopegelonline==0.0.5"]
+  "requirements": ["aiopegelonline==0.0.6"]
 }

homeassistant/components/pegel_online/model.py

@@ -1,11 +0,0 @@
-"""Models for PEGELONLINE."""
-
-from typing import TypedDict
-
-from aiopegelonline import CurrentMeasurement
-
-
-class PegelOnlineData(TypedDict):
-    """TypedDict for PEGELONLINE Coordinator Data."""
-
-    water_level: CurrentMeasurement

homeassistant/components/pegel_online/sensor.py

@@ -1,10 +1,12 @@
 """PEGELONLINE sensor entities."""
 from __future__ import annotations
 
-from collections.abc import Callable
 from dataclasses import dataclass
 
+from aiopegelonline.models import CurrentMeasurement
+
 from homeassistant.components.sensor import (
+    SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
     SensorStateClass,
@@ -17,15 +19,13 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from .const import DOMAIN
 from .coordinator import PegelOnlineDataUpdateCoordinator
 from .entity import PegelOnlineEntity
-from .model import PegelOnlineData
 
 
 @dataclass
 class PegelOnlineRequiredKeysMixin:
     """Mixin for required keys."""
 
-    fn_native_unit: Callable[[PegelOnlineData], str]
+    measurement_key: str
-    fn_native_value: Callable[[PegelOnlineData], float]
 
 
 @dataclass
@@ -36,14 +36,71 @@ class PegelOnlineSensorEntityDescription(
 
 
 SENSORS: tuple[PegelOnlineSensorEntityDescription, ...] = (
+    PegelOnlineSensorEntityDescription(
+        key="air_temperature",
+        translation_key="air_temperature",
+        measurement_key="air_temperature",
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.TEMPERATURE,
+        icon="mdi:thermometer-lines",
+        entity_registry_enabled_default=False,
+    ),
+    PegelOnlineSensorEntityDescription(
+        key="clearance_height",
+        translation_key="clearance_height",
+        measurement_key="clearance_height",
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.DISTANCE,
+        icon="mdi:bridge",
+    ),
+    PegelOnlineSensorEntityDescription(
+        key="oxygen_level",
+        translation_key="oxygen_level",
+        measurement_key="oxygen_level",
+        state_class=SensorStateClass.MEASUREMENT,
+        icon="mdi:water-opacity",
+        entity_registry_enabled_default=False,
+    ),
+    PegelOnlineSensorEntityDescription(
+        key="ph_value",
+        measurement_key="ph_value",
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.PH,
+        entity_registry_enabled_default=False,
+    ),
+    PegelOnlineSensorEntityDescription(
+        key="water_speed",
+        translation_key="water_speed",
+        measurement_key="water_speed",
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.SPEED,
+        icon="mdi:waves-arrow-right",
+        entity_registry_enabled_default=False,
+    ),
+    PegelOnlineSensorEntityDescription(
+        key="water_flow",
+        translation_key="water_flow",
+        measurement_key="water_flow",
+        state_class=SensorStateClass.MEASUREMENT,
+        icon="mdi:waves",
+        entity_registry_enabled_default=False,
+    ),
     PegelOnlineSensorEntityDescription(
         key="water_level",
         translation_key="water_level",
+        measurement_key="water_level",
         state_class=SensorStateClass.MEASUREMENT,
-        fn_native_unit=lambda data: data["water_level"].uom,
-        fn_native_value=lambda data: data["water_level"].value,
         icon="mdi:waves-arrow-up",
     ),
+    PegelOnlineSensorEntityDescription(
+        key="water_temperature",
+        translation_key="water_temperature",
+        measurement_key="water_temperature",
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.TEMPERATURE,
+        icon="mdi:thermometer-water",
+        entity_registry_enabled_default=False,
+    ),
 )
 
 
@@ -51,9 +108,14 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the PEGELONLINE sensor."""
-    coordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: PegelOnlineDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+
     async_add_entities(
-        [PegelOnlineSensor(coordinator, description) for description in SENSORS]
+        [
+            PegelOnlineSensor(coordinator, description)
+            for description in SENSORS
+            if getattr(coordinator.data, description.measurement_key) is not None
+        ]
     )
 
 
@@ -71,9 +133,9 @@ class PegelOnlineSensor(PegelOnlineEntity, SensorEntity):
         super().__init__(coordinator)
         self.entity_description = description
         self._attr_unique_id = f"{self.station.uuid}_{description.key}"
-        self._attr_native_unit_of_measurement = self.entity_description.fn_native_unit(
+
-            coordinator.data
+        if description.device_class != SensorDeviceClass.PH:
-        )
+            self._attr_native_unit_of_measurement = self.measurement.uom
 
         if self.station.latitude and self.station.longitude:
             self._attr_extra_state_attributes.update(
@@ -83,7 +145,12 @@ class PegelOnlineSensor(PegelOnlineEntity, SensorEntity):
                 }
             )
 
+    @property
+    def measurement(self) -> CurrentMeasurement:
+        """Return the measurement data of the entity."""
+        return getattr(self.coordinator.data, self.entity_description.measurement_key)
+
     @property
     def native_value(self) -> float:
         """Return the state of the device."""
-        return self.entity_description.fn_native_value(self.coordinator.data)
+        return self.measurement.value

homeassistant/components/pegel_online/strings.json

@@ -26,8 +26,26 @@
   },
   "entity": {
     "sensor": {
+      "air_temperature": {
+        "name": "Air temperature"
+      },
+      "clearance_height": {
+        "name": "Clearance height"
+      },
+      "oxygen_level": {
+        "name": "Oxygen level"
+      },
+      "water_speed": {
+        "name": "Water flow speed"
+      },
+      "water_flow": {
+        "name": "Water volume flow"
+      },
       "water_level": {
         "name": "Water level"
+      },
+      "water_temperature": {
+        "name": "Water temperature"
       }
     }
   }

requirements_all.txt

@@ -304,7 +304,7 @@ aiooncue==0.3.5
 aioopenexchangerates==0.4.0
 
 # homeassistant.components.pegel_online
-aiopegelonline==0.0.5
+aiopegelonline==0.0.6
 
 # homeassistant.components.acmeda
 aiopulse==0.4.3

requirements_test_all.txt

@@ -279,7 +279,7 @@ aiooncue==0.3.5
 aioopenexchangerates==0.4.0
 
 # homeassistant.components.pegel_online
-aiopegelonline==0.0.5
+aiopegelonline==0.0.6
 
 # homeassistant.components.acmeda
 aiopulse==0.4.3

tests/components/pegel_online/__init__.py

@@ -8,13 +8,13 @@ class PegelOnlineMock:
         self,
         nearby_stations=None,
         station_details=None,
-        station_measurement=None,
+        station_measurements=None,
         side_effect=None,
     ) -> None:
         """Init the mock."""
         self.nearby_stations = nearby_stations
         self.station_details = station_details
-        self.station_measurement = station_measurement
+        self.station_measurements = station_measurements
         self.side_effect = side_effect
 
     async def async_get_nearby_stations(self, *args):
@@ -29,11 +29,11 @@ class PegelOnlineMock:
             raise self.side_effect
         return self.station_details
 
-    async def async_get_station_measurement(self, *args):
+    async def async_get_station_measurements(self, *args):
-        """Mock async_get_station_measurement."""
+        """Mock async_get_station_measurements."""
         if self.side_effect:
             raise self.side_effect
-        return self.station_measurement
+        return self.station_measurements
 
     def override_side_effect(self, side_effect):
         """Override the side_effect."""

tests/components/pegel_online/const.py

@@ -0,0 +1,203 @@
+"""Constants for pegel_online tests."""
+
+from aiopegelonline.models import Station, StationMeasurements
+
+from homeassistant.components.pegel_online.const import CONF_STATION
+
+MOCK_STATION_DETAILS_MEISSEN = Station(
+    {
+        "uuid": "85d686f1-xxxx-xxxx-xxxx-3207b50901a7",
+        "number": "501060",
+        "shortname": "MEISSEN",
+        "longname": "MEISSEN",
+        "km": 82.2,
+        "agency": "STANDORT DRESDEN",
+        "longitude": 13.475467710324812,
+        "latitude": 51.16440557554545,
+        "water": {"shortname": "ELBE", "longname": "ELBE"},
+    }
+)
+
+MOCK_STATION_DETAILS_DRESDEN = Station(
+    {
+        "uuid": "70272185-xxxx-xxxx-xxxx-43bea330dcae",
+        "number": "501060",
+        "shortname": "DRESDEN",
+        "longname": "DRESDEN",
+        "km": 55.63,
+        "agency": "STANDORT DRESDEN",
+        "longitude": 13.738831783620384,
+        "latitude": 51.054459765598125,
+        "water": {"shortname": "ELBE", "longname": "ELBE"},
+    }
+)
+MOCK_CONFIG_ENTRY_DATA_DRESDEN = {CONF_STATION: "70272185-xxxx-xxxx-xxxx-43bea330dcae"}
+MOCK_STATION_MEASUREMENT_DRESDEN = StationMeasurements(
+    [
+        {
+            "shortname": "W",
+            "longname": "WASSERSTAND ROHDATEN",
+            "unit": "cm",
+            "equidistance": 15,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T21:15:00+02:00",
+                "value": 62,
+                "stateMnwMhw": "low",
+                "stateNswHsw": "normal",
+            },
+            "gaugeZero": {
+                "unit": "m. ü. NHN",
+                "value": 102.7,
+                "validFrom": "2019-11-01",
+            },
+        },
+        {
+            "shortname": "Q",
+            "longname": "ABFLUSS_ROHDATEN",
+            "unit": "m³/s",
+            "equidistance": 15,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T06:00:00+02:00",
+                "value": 88.4,
+            },
+        },
+    ]
+)
+
+MOCK_STATION_DETAILS_HANAU_BRIDGE = Station(
+    {
+        "uuid": "07374faf-xxxx-xxxx-xxxx-adc0e0784c4b",
+        "number": "24700347",
+        "shortname": "HANAU BRÜCKE DFH",
+        "longname": "HANAU BRÜCKE DFH",
+        "km": 56.398,
+        "agency": "ASCHAFFENBURG",
+        "water": {"shortname": "MAIN", "longname": "MAIN"},
+    }
+)
+MOCK_CONFIG_ENTRY_DATA_HANAU_BRIDGE = {
+    CONF_STATION: "07374faf-xxxx-xxxx-xxxx-adc0e0784c4b"
+}
+MOCK_STATION_MEASUREMENT_HANAU_BRIDGE = StationMeasurements(
+    [
+        {
+            "shortname": "DFH",
+            "longname": "DURCHFAHRTSHÖHE",
+            "unit": "cm",
+            "equidistance": 15,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:45:00+02:00",
+                "value": 715,
+            },
+            "gaugeZero": {
+                "unit": "m. ü. NHN",
+                "value": 106.501,
+                "validFrom": "2019-11-01",
+            },
+        }
+    ]
+)
+
+
+MOCK_STATION_DETAILS_WUERZBURG = Station(
+    {
+        "uuid": "915d76e1-xxxx-xxxx-xxxx-4d144cd771cc",
+        "number": "24300600",
+        "shortname": "WÜRZBURG",
+        "longname": "WÜRZBURG",
+        "km": 251.97,
+        "agency": "SCHWEINFURT",
+        "longitude": 9.925968763247354,
+        "latitude": 49.79620901036012,
+        "water": {"shortname": "MAIN", "longname": "MAIN"},
+    }
+)
+MOCK_CONFIG_ENTRY_DATA_WUERZBURG = {
+    CONF_STATION: "915d76e1-xxxx-xxxx-xxxx-4d144cd771cc"
+}
+MOCK_STATION_MEASUREMENT_WUERZBURG = StationMeasurements(
+    [
+        {
+            "shortname": "W",
+            "longname": "WASSERSTAND ROHDATEN",
+            "unit": "cm",
+            "equidistance": 15,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:15:00+02:00",
+                "value": 159,
+                "stateMnwMhw": "normal",
+                "stateNswHsw": "normal",
+            },
+            "gaugeZero": {
+                "unit": "m. ü. NHN",
+                "value": 164.511,
+                "validFrom": "2019-11-01",
+            },
+        },
+        {
+            "shortname": "LT",
+            "longname": "LUFTTEMPERATUR",
+            "unit": "°C",
+            "equidistance": 60,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:00:00+02:00",
+                "value": 21.2,
+            },
+        },
+        {
+            "shortname": "WT",
+            "longname": "WASSERTEMPERATUR",
+            "unit": "°C",
+            "equidistance": 60,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:00:00+02:00",
+                "value": 22.1,
+            },
+        },
+        {
+            "shortname": "VA",
+            "longname": "FLIESSGESCHWINDIGKEIT",
+            "unit": "m/s",
+            "equidistance": 15,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:15:00+02:00",
+                "value": 0.58,
+            },
+        },
+        {
+            "shortname": "O2",
+            "longname": "SAUERSTOFFGEHALT",
+            "unit": "mg/l",
+            "equidistance": 60,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:00:00+02:00",
+                "value": 8.4,
+            },
+        },
+        {
+            "shortname": "PH",
+            "longname": "PH-WERT",
+            "unit": "--",
+            "equidistance": 60,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:00:00+02:00",
+                "value": 8.1,
+            },
+        },
+        {
+            "shortname": "Q",
+            "longname": "ABFLUSS",
+            "unit": "m³/s",
+            "equidistance": 15,
+            "currentMeasurement": {
+                "timestamp": "2023-07-26T19:00:00+02:00",
+                "value": 102,
+            },
+        },
+    ]
+)
+
+MOCK_NEARBY_STATIONS = {
+    "70272185-xxxx-xxxx-xxxx-43bea330dcae": MOCK_STATION_DETAILS_DRESDEN,
+    "85d686f1-xxxx-xxxx-xxxx-3207b50901a7": MOCK_STATION_DETAILS_MEISSEN,
+}

tests/components/pegel_online/test_config_flow.py

@@ -2,7 +2,6 @@
 from unittest.mock import patch
 
 from aiohttp.client_exceptions import ClientError
-from aiopegelonline import Station
 
 from homeassistant.components.pegel_online.const import (
     CONF_STATION,
@@ -19,6 +18,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
 from . import PegelOnlineMock
+from .const import MOCK_CONFIG_ENTRY_DATA_DRESDEN, MOCK_NEARBY_STATIONS
 
 from tests.common import MockConfigEntry
 
@@ -27,38 +27,7 @@ MOCK_USER_DATA_STEP1 = {
     CONF_RADIUS: 25,
 }
 
-MOCK_USER_DATA_STEP2 = {CONF_STATION: "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"}
+MOCK_USER_DATA_STEP2 = {CONF_STATION: "70272185-xxxx-xxxx-xxxx-43bea330dcae"}
-
-MOCK_CONFIG_ENTRY_DATA = {CONF_STATION: "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"}
-
-MOCK_NEARBY_STATIONS = {
-    "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8": Station(
-        {
-            "uuid": "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8",
-            "number": "501060",
-            "shortname": "DRESDEN",
-            "longname": "DRESDEN",
-            "km": 55.63,
-            "agency": "STANDORT DRESDEN",
-            "longitude": 13.738831783620384,
-            "latitude": 51.054459765598125,
-            "water": {"shortname": "ELBE", "longname": "ELBE"},
-        }
-    ),
-    "85d686f1-xxxx-xxxx-xxxx-3207b50901a7": Station(
-        {
-            "uuid": "85d686f1-xxxx-xxxx-xxxx-3207b50901a7",
-            "number": "501060",
-            "shortname": "MEISSEN",
-            "longname": "MEISSEN",
-            "km": 82.2,
-            "agency": "STANDORT DRESDEN",
-            "longitude": 13.475467710324812,
-            "latitude": 51.16440557554545,
-            "water": {"shortname": "ELBE", "longname": "ELBE"},
-        }
-    ),
-}
 
 
 async def test_user(hass: HomeAssistant) -> None:
@@ -85,7 +54,7 @@ async def test_user(hass: HomeAssistant) -> None:
             result["flow_id"], user_input=MOCK_USER_DATA_STEP2
         )
         assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["data"][CONF_STATION] == "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"
+        assert result["data"][CONF_STATION] == "70272185-xxxx-xxxx-xxxx-43bea330dcae"
         assert result["title"] == "DRESDEN ELBE"
 
         await hass.async_block_till_done()
@@ -97,8 +66,8 @@ async def test_user_already_configured(hass: HomeAssistant) -> None:
     """Test starting a flow by user with an already configured statioon."""
     mock_config = MockConfigEntry(
         domain=DOMAIN,
-        data=MOCK_CONFIG_ENTRY_DATA,
+        data=MOCK_CONFIG_ENTRY_DATA_DRESDEN,
-        unique_id=MOCK_CONFIG_ENTRY_DATA[CONF_STATION],
+        unique_id=MOCK_CONFIG_ENTRY_DATA_DRESDEN[CONF_STATION],
     )
     mock_config.add_to_hass(hass)
 
@@ -159,7 +128,7 @@ async def test_connection_error(hass: HomeAssistant) -> None:
             result["flow_id"], user_input=MOCK_USER_DATA_STEP2
         )
         assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["data"][CONF_STATION] == "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"
+        assert result["data"][CONF_STATION] == "70272185-xxxx-xxxx-xxxx-43bea330dcae"
         assert result["title"] == "DRESDEN ELBE"
 
         await hass.async_block_till_done()
@@ -201,7 +170,7 @@ async def test_user_no_stations(hass: HomeAssistant) -> None:
             result["flow_id"], user_input=MOCK_USER_DATA_STEP2
         )
         assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["data"][CONF_STATION] == "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"
+        assert result["data"][CONF_STATION] == "70272185-xxxx-xxxx-xxxx-43bea330dcae"
         assert result["title"] == "DRESDEN ELBE"
 
         await hass.async_block_till_done()

tests/components/pegel_online/test_init.py

@@ -2,7 +2,6 @@
 from unittest.mock import patch
 
 from aiohttp.client_exceptions import ClientError
-from aiopegelonline import CurrentMeasurement, Station
 
 from homeassistant.components.pegel_online.const import (
     CONF_STATION,
@@ -14,39 +13,27 @@ from homeassistant.core import HomeAssistant
 from homeassistant.util import utcnow
 
 from . import PegelOnlineMock
+from .const import (
+    MOCK_CONFIG_ENTRY_DATA_DRESDEN,
+    MOCK_STATION_DETAILS_DRESDEN,
+    MOCK_STATION_MEASUREMENT_DRESDEN,
+)
 
 from tests.common import MockConfigEntry, async_fire_time_changed
 
-MOCK_CONFIG_ENTRY_DATA = {CONF_STATION: "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"}
-
-MOCK_STATION_DETAILS = Station(
-    {
-        "uuid": "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8",
-        "number": "501060",
-        "shortname": "DRESDEN",
-        "longname": "DRESDEN",
-        "km": 55.63,
-        "agency": "STANDORT DRESDEN",
-        "longitude": 13.738831783620384,
-        "latitude": 51.054459765598125,
-        "water": {"shortname": "ELBE", "longname": "ELBE"},
-    }
-)
-MOCK_STATION_MEASUREMENT = CurrentMeasurement("cm", 56)
-
 
 async def test_update_error(hass: HomeAssistant) -> None:
     """Tests error during update entity."""
     entry = MockConfigEntry(
         domain=DOMAIN,
-        data=MOCK_CONFIG_ENTRY_DATA,
+        data=MOCK_CONFIG_ENTRY_DATA_DRESDEN,
-        unique_id=MOCK_CONFIG_ENTRY_DATA[CONF_STATION],
+        unique_id=MOCK_CONFIG_ENTRY_DATA_DRESDEN[CONF_STATION],
     )
     entry.add_to_hass(hass)
     with patch("homeassistant.components.pegel_online.PegelOnline") as pegelonline:
         pegelonline.return_value = PegelOnlineMock(
-            station_details=MOCK_STATION_DETAILS,
+            station_details=MOCK_STATION_DETAILS_DRESDEN,
-            station_measurement=MOCK_STATION_MEASUREMENT,
+            station_measurements=MOCK_STATION_MEASUREMENT_DRESDEN,
         )
         assert await hass.config_entries.async_setup(entry.entry_id)
 

tests/components/pegel_online/test_sensor.py

@@ -1,53 +1,141 @@
 """Test pegel_online component."""
 from unittest.mock import patch
 
-from aiopegelonline import CurrentMeasurement, Station
+from aiopegelonline.models import Station, StationMeasurements
+import pytest
 
 from homeassistant.components.pegel_online.const import CONF_STATION, DOMAIN
-from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE
+from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE, ATTR_UNIT_OF_MEASUREMENT
 from homeassistant.core import HomeAssistant
 
 from . import PegelOnlineMock
+from .const import (
+    MOCK_CONFIG_ENTRY_DATA_DRESDEN,
+    MOCK_CONFIG_ENTRY_DATA_HANAU_BRIDGE,
+    MOCK_CONFIG_ENTRY_DATA_WUERZBURG,
+    MOCK_STATION_DETAILS_DRESDEN,
+    MOCK_STATION_DETAILS_HANAU_BRIDGE,
+    MOCK_STATION_DETAILS_WUERZBURG,
+    MOCK_STATION_MEASUREMENT_DRESDEN,
+    MOCK_STATION_MEASUREMENT_HANAU_BRIDGE,
+    MOCK_STATION_MEASUREMENT_WUERZBURG,
+)
 
 from tests.common import MockConfigEntry
 
-MOCK_CONFIG_ENTRY_DATA = {CONF_STATION: "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8"}
 
-MOCK_STATION_DETAILS = Station(
+@pytest.mark.parametrize(
-    {
+    (
-        "uuid": "3bcd61da-xxxx-xxxx-xxxx-19d5523a7ae8",
+        "mock_config_entry_data",
-        "number": "501060",
+        "mock_station_details",
-        "shortname": "DRESDEN",
+        "mock_station_measurement",
-        "longname": "DRESDEN",
+        "expected_states",
-        "km": 55.63,
+    ),
-        "agency": "STANDORT DRESDEN",
+    [
-        "longitude": 13.738831783620384,
+        (
-        "latitude": 51.054459765598125,
+            MOCK_CONFIG_ENTRY_DATA_DRESDEN,
-        "water": {"shortname": "ELBE", "longname": "ELBE"},
+            MOCK_STATION_DETAILS_DRESDEN,
-    }
+            MOCK_STATION_MEASUREMENT_DRESDEN,
+            {
+                "sensor.dresden_elbe_water_volume_flow": (
+                    "DRESDEN ELBE Water volume flow",
+                    "88.4",
+                    "m³/s",
+                ),
+                "sensor.dresden_elbe_water_level": (
+                    "DRESDEN ELBE Water level",
+                    "62",
+                    "cm",
+                ),
+            },
+        ),
+        (
+            MOCK_CONFIG_ENTRY_DATA_HANAU_BRIDGE,
+            MOCK_STATION_DETAILS_HANAU_BRIDGE,
+            MOCK_STATION_MEASUREMENT_HANAU_BRIDGE,
+            {
+                "sensor.hanau_brucke_dfh_main_clearance_height": (
+                    "HANAU BRÜCKE DFH MAIN Clearance height",
+                    "715",
+                    "cm",
+                ),
+            },
+        ),
+        (
+            MOCK_CONFIG_ENTRY_DATA_WUERZBURG,
+            MOCK_STATION_DETAILS_WUERZBURG,
+            MOCK_STATION_MEASUREMENT_WUERZBURG,
+            {
+                "sensor.wurzburg_main_air_temperature": (
+                    "WÜRZBURG MAIN Air temperature",
+                    "21.2",
+                    "°C",
+                ),
+                "sensor.wurzburg_main_oxygen_level": (
+                    "WÜRZBURG MAIN Oxygen level",
+                    "8.4",
+                    "mg/l",
+                ),
+                "sensor.wurzburg_main_ph": (
+                    "WÜRZBURG MAIN pH",
+                    "8.1",
+                    None,
+                ),
+                "sensor.wurzburg_main_water_flow_speed": (
+                    "WÜRZBURG MAIN Water flow speed",
+                    "0.58",
+                    "m/s",
+                ),
+                "sensor.wurzburg_main_water_volume_flow": (
+                    "WÜRZBURG MAIN Water volume flow",
+                    "102",
+                    "m³/s",
+                ),
+                "sensor.wurzburg_main_water_level": (
+                    "WÜRZBURG MAIN Water level",
+                    "159",
+                    "cm",
+                ),
+                "sensor.wurzburg_main_water_temperature": (
+                    "WÜRZBURG MAIN Water temperature",
+                    "22.1",
+                    "°C",
+                ),
+            },
+        ),
+    ],
 )
-MOCK_STATION_MEASUREMENT = CurrentMeasurement("cm", 56)
+async def test_sensor(
-
+    hass: HomeAssistant,
-
+    mock_config_entry_data: dict,
-async def test_sensor(hass: HomeAssistant) -> None:
+    mock_station_details: Station,
+    mock_station_measurement: StationMeasurements,
+    expected_states: dict,
+    entity_registry_enabled_by_default: None,
+) -> None:
     """Tests sensor entity."""
     entry = MockConfigEntry(
         domain=DOMAIN,
-        data=MOCK_CONFIG_ENTRY_DATA,
+        data=mock_config_entry_data,
-        unique_id=MOCK_CONFIG_ENTRY_DATA[CONF_STATION],
+        unique_id=mock_config_entry_data[CONF_STATION],
     )
     entry.add_to_hass(hass)
     with patch("homeassistant.components.pegel_online.PegelOnline") as pegelonline:
         pegelonline.return_value = PegelOnlineMock(
-            station_details=MOCK_STATION_DETAILS,
+            station_details=mock_station_details,
-            station_measurement=MOCK_STATION_MEASUREMENT,
+            station_measurements=mock_station_measurement,
         )
         assert await hass.config_entries.async_setup(entry.entry_id)
 
     await hass.async_block_till_done()
 
-    state = hass.states.get("sensor.dresden_elbe_water_level")
+    assert len(hass.states.async_all()) == len(expected_states)
-    assert state.name == "DRESDEN ELBE Water level"
+
-    assert state.state == "56"
+    for state_name, state_data in expected_states.items():
-    assert state.attributes[ATTR_LATITUDE] == 51.054459765598125
+        state = hass.states.get(state_name)
-    assert state.attributes[ATTR_LONGITUDE] == 13.738831783620384
+        assert state.name == state_data[0]
+        assert state.state == state_data[1]
+        assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == state_data[2]
+        if mock_station_details.latitude is not None:
+            assert state.attributes[ATTR_LATITUDE] == mock_station_details.latitude
+            assert state.attributes[ATTR_LONGITUDE] == mock_station_details.longitude