Introduce base entity in streamlabs water

homeassistant/components/streamlabswater/binary_sensor.py

@@ -5,13 +5,10 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import StreamlabsCoordinator
 from .const import DOMAIN
-from .coordinator import StreamlabsData
-
-NAME_AWAY_MODE = "Water Away Mode"
+from .entity import StreamlabsWaterEntity
 
 
 async def async_setup_entry(
@@ -22,31 +19,19 @@ async def async_setup_entry(
     """Set up Streamlabs water binary sensor from a config entry."""
     coordinator = hass.data[DOMAIN][entry.entry_id]
 
-    entities = []
-
-    for location_id in coordinator.data:
-        entities.append(StreamlabsAwayMode(coordinator, location_id))
+    async_add_entities(
+        StreamlabsAwayMode(coordinator, location_id) for location_id in coordinator.data
+    )
 
-    async_add_entities(entities)
 
-
-class StreamlabsAwayMode(CoordinatorEntity[StreamlabsCoordinator], BinarySensorEntity):
+class StreamlabsAwayMode(StreamlabsWaterEntity, BinarySensorEntity):
     """Monitor the away mode state."""
 
+    _attr_translation_key = "away_mode"
+
     def __init__(self, coordinator: StreamlabsCoordinator, location_id: str) -> None:
         """Initialize the away mode device."""
-        super().__init__(coordinator)
-        self._location_id = location_id
-
-    @property
-    def location_data(self) -> StreamlabsData:
-        """Returns the data object."""
-        return self.coordinator.data[self._location_id]
-
-    @property
-    def name(self) -> str:
-        """Return the name for away mode."""
-        return f"{self.location_data.name} {NAME_AWAY_MODE}"
+        super().__init__(coordinator, location_id, "away_mode")
 
     @property
     def is_on(self) -> bool:

homeassistant/components/streamlabswater/entity.py

@@ -0,0 +1,31 @@
+"""Base entity for Streamlabs integration."""
+from homeassistant.core import DOMAIN
+from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .coordinator import StreamlabsCoordinator, StreamlabsData
+
+
+class StreamlabsWaterEntity(CoordinatorEntity[StreamlabsCoordinator]):
+    """Defines a base Streamlabs entity."""
+
+    _attr_has_entity_name = True
+
+    def __init__(
+        self,
+        coordinator: StreamlabsCoordinator,
+        location_id: str,
+        key: str,
+    ) -> None:
+        """Initialize the Streamlabs entity."""
+        super().__init__(coordinator)
+        self._location_id = location_id
+        self._attr_unique_id = f"{location_id}-{key}"
+        self._attr_device_info = DeviceInfo(
+            identifiers={(DOMAIN, location_id)}, name=self.location_data.name
+        )
+
+    @property
+    def location_data(self) -> StreamlabsData:
+        """Returns the data object."""
+        return self.coordinator.data[self._location_id]

homeassistant/components/streamlabswater/sensor.py

@@ -12,14 +12,13 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfVolume
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import StreamlabsCoordinator
 from .const import DOMAIN
 from .coordinator import StreamlabsData
+from .entity import StreamlabsWaterEntity
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -69,11 +68,9 @@ async def async_setup_entry(
     )
 
 
-class StreamLabsSensor(CoordinatorEntity[StreamlabsCoordinator], SensorEntity):
+class StreamLabsSensor(StreamlabsWaterEntity, SensorEntity):
     """Monitors the daily water usage."""
 
-    _attr_has_entity_name = True
-
     entity_description: StreamlabsWaterSensorEntityDescription
 
     def __init__(
@@ -83,18 +80,8 @@ def __init__(
         entity_description: StreamlabsWaterSensorEntityDescription,
     ) -> None:
         """Initialize the daily water usage device."""
-        super().__init__(coordinator)
-        self._location_id = location_id
-        self._attr_unique_id = f"{location_id}-{entity_description.key}"
+        super().__init__(coordinator, location_id, entity_description.key)
         self.entity_description = entity_description
-        self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, location_id)}, name=self.location_data.name
-        )
-
-    @property
-    def location_data(self) -> StreamlabsData:
-        """Returns the data object."""
-        return self.coordinator.data[self._location_id]
 
     @property
     def native_value(self) -> StateType:

homeassistant/components/streamlabswater/strings.json

@@ -32,6 +32,11 @@
     }
   },
   "entity": {
+    "binary_sensor": {
+      "away_mode": {
+        "name": "Away mode"
+      }
+    },
     "sensor": {
       "daily_usage": {
         "name": "Daily usage"

tests/components/streamlabswater/snapshots/test_binary_sensor.ambr

@@ -0,0 +1,44 @@
+# serializer version: 1
+# name: test_all_entities[binary_sensor.water_monitor_away_mode-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'binary_sensor',
+    'entity_category': None,
+    'entity_id': 'binary_sensor.water_monitor_away_mode',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'Away mode',
+    'platform': 'streamlabswater',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'away_mode',
+    'unique_id': '945e7c52-854a-41e1-8524-50c6993277e1-away_mode',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_all_entities[binary_sensor.water_monitor_away_mode-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Water Monitor Away mode',
+    }),
+    'context': <ANY>,
+    'entity_id': 'binary_sensor.water_monitor_away_mode',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'off',
+  })
+# ---

tests/components/streamlabswater/test_binary_sensor.py

@@ -0,0 +1,35 @@
+"""Tests for the Streamlabs Water binary sensor platform."""
+from unittest.mock import AsyncMock, patch
+
+from syrupy import SnapshotAssertion
+
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+
+from tests.common import MockConfigEntry
+from tests.components.streamlabswater import setup_integration
+
+
+async def test_all_entities(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    streamlabswater: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test all entities."""
+    with patch(
+        "homeassistant.components.streamlabswater.PLATFORMS", [Platform.BINARY_SENSOR]
+    ):
+        await setup_integration(hass, mock_config_entry)
+        entity_entries = er.async_entries_for_config_entry(
+            entity_registry, mock_config_entry.entry_id
+        )
+
+        assert entity_entries
+        for entity_entry in entity_entries:
+            assert entity_entry == snapshot(name=f"{entity_entry.entity_id}-entry")
+            assert hass.states.get(entity_entry.entity_id) == snapshot(
+                name=f"{entity_entry.entity_id}-state"
+            )