From 52aec885ea3dc2651c4c095c5432747253e842de Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 23 Sep 2024 12:32:18 +0200
Subject: [PATCH] Move nibe_heatpump base entity to separate module (#126498)

---
 .../components/nibe_heatpump/__init__.py      |  4 +-
 .../components/nibe_heatpump/binary_sensor.py |  7 +--
 .../components/nibe_heatpump/button.py        |  8 +--
 .../components/nibe_heatpump/climate.py       |  8 +--
 .../components/nibe_heatpump/coordinator.py   | 49 +-----------------
 .../components/nibe_heatpump/entity.py        | 50 +++++++++++++++++++
 .../components/nibe_heatpump/number.py        |  7 +--
 .../components/nibe_heatpump/select.py        |  7 +--
 .../components/nibe_heatpump/sensor.py        |  7 +--
 .../components/nibe_heatpump/switch.py        |  7 +--
 .../components/nibe_heatpump/water_heater.py  |  8 +--
 11 files changed, 86 insertions(+), 76 deletions(-)
 create mode 100644 homeassistant/components/nibe_heatpump/entity.py

diff --git a/homeassistant/components/nibe_heatpump/__init__.py b/homeassistant/components/nibe_heatpump/__init__.py
index fbb49351e0e..b3ceb00a834 100644
--- a/homeassistant/components/nibe_heatpump/__init__.py
+++ b/homeassistant/components/nibe_heatpump/__init__.py
@@ -30,7 +30,7 @@ from .const import (
     CONF_WORD_SWAP,
     DOMAIN,
 )
-from .coordinator import Coordinator
+from .coordinator import CoilCoordinator
 
 PLATFORMS: list[Platform] = [
     Platform.BINARY_SENSOR,
@@ -81,7 +81,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_stop)
     )
 
-    coordinator = Coordinator(hass, heatpump, connection)
+    coordinator = CoilCoordinator(hass, heatpump, connection)
 
     data = hass.data.setdefault(DOMAIN, {})
     data[entry.entry_id] = coordinator
diff --git a/homeassistant/components/nibe_heatpump/binary_sensor.py b/homeassistant/components/nibe_heatpump/binary_sensor.py
index 035a4a23a08..0cb16bf4485 100644
--- a/homeassistant/components/nibe_heatpump/binary_sensor.py
+++ b/homeassistant/components/nibe_heatpump/binary_sensor.py
@@ -11,7 +11,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import CoilEntity, Coordinator
+from .coordinator import CoilCoordinator
+from .entity import CoilEntity
 
 
 async def async_setup_entry(
@@ -21,7 +22,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     async_add_entities(
         BinarySensor(coordinator, coil)
@@ -35,7 +36,7 @@ class BinarySensor(CoilEntity, BinarySensorEntity):
 
     _attr_entity_category = EntityCategory.DIAGNOSTIC
 
-    def __init__(self, coordinator: Coordinator, coil: Coil) -> None:
+    def __init__(self, coordinator: CoilCoordinator, coil: Coil) -> None:
         """Initialize entity."""
         super().__init__(coordinator, coil, ENTITY_ID_FORMAT)
 
diff --git a/homeassistant/components/nibe_heatpump/button.py b/homeassistant/components/nibe_heatpump/button.py
index 0c3122805e1..df8ceef6479 100644
--- a/homeassistant/components/nibe_heatpump/button.py
+++ b/homeassistant/components/nibe_heatpump/button.py
@@ -13,7 +13,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN, LOGGER
-from .coordinator import Coordinator
+from .coordinator import CoilCoordinator
 
 
 async def async_setup_entry(
@@ -23,7 +23,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     def reset_buttons():
         if unit := UNIT_COILGROUPS.get(coordinator.series, {}).get("main"):
@@ -35,13 +35,13 @@ async def async_setup_entry(
     async_add_entities(reset_buttons())
 
 
-class NibeAlarmResetButton(CoordinatorEntity[Coordinator], ButtonEntity):
+class NibeAlarmResetButton(CoordinatorEntity[CoilCoordinator], ButtonEntity):
     """Sensor entity."""
 
     _attr_has_entity_name = True
     _attr_entity_category = EntityCategory.DIAGNOSTIC
 
-    def __init__(self, coordinator: Coordinator, unit: UnitCoilGroup) -> None:
+    def __init__(self, coordinator: CoilCoordinator, unit: UnitCoilGroup) -> None:
         """Initialize entity."""
         self._reset_coil = coordinator.heatpump.get_coil_by_address(unit.alarm_reset)
         self._alarm_coil = coordinator.heatpump.get_coil_by_address(unit.alarm)
diff --git a/homeassistant/components/nibe_heatpump/climate.py b/homeassistant/components/nibe_heatpump/climate.py
index d933d5a5ab0..f89d6ec29a9 100644
--- a/homeassistant/components/nibe_heatpump/climate.py
+++ b/homeassistant/components/nibe_heatpump/climate.py
@@ -38,7 +38,7 @@ from .const import (
     VALUES_PRIORITY_COOLING,
     VALUES_PRIORITY_HEATING,
 )
-from .coordinator import Coordinator
+from .coordinator import CoilCoordinator
 
 
 async def async_setup_entry(
@@ -48,7 +48,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     main_unit = UNIT_COILGROUPS[coordinator.series]["main"]
 
@@ -62,7 +62,7 @@ async def async_setup_entry(
     async_add_entities(climate_systems())
 
 
-class NibeClimateEntity(CoordinatorEntity[Coordinator], ClimateEntity):
+class NibeClimateEntity(CoordinatorEntity[CoilCoordinator], ClimateEntity):
     """Climate entity."""
 
     _attr_entity_category = None
@@ -78,7 +78,7 @@ class NibeClimateEntity(CoordinatorEntity[Coordinator], ClimateEntity):
 
     def __init__(
         self,
-        coordinator: Coordinator,
+        coordinator: CoilCoordinator,
         key: str,
         unit: UnitCoilGroup,
         climate: ClimateCoilGroup,
diff --git a/homeassistant/components/nibe_heatpump/coordinator.py b/homeassistant/components/nibe_heatpump/coordinator.py
index 0f1fabe4249..2c19703549a 100644
--- a/homeassistant/components/nibe_heatpump/coordinator.py
+++ b/homeassistant/components/nibe_heatpump/coordinator.py
@@ -17,12 +17,7 @@ from nibe.heatpump import HeatPump, Series
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.helpers.device_registry import DeviceInfo
-from homeassistant.helpers.entity import async_generate_entity_id
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-    UpdateFailed,
-)
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN, LOGGER
 
@@ -68,7 +63,7 @@ class ContextCoordinator[_DataTypeT, _ContextTypeT](DataUpdateCoordinator[_DataT
         return release_update
 
 
-class Coordinator(ContextCoordinator[dict[int, CoilData], int]):
+class CoilCoordinator(ContextCoordinator[dict[int, CoilData], int]):
     """Update coordinator for nibe heat pumps."""
 
     config_entry: ConfigEntry
@@ -188,43 +183,3 @@ class Coordinator(ContextCoordinator[dict[int, CoilData], int]):
             self.task.cancel()
             await asyncio.wait((self.task,))
         await self.connection.stop()
-
-
-class CoilEntity(CoordinatorEntity[Coordinator]):
-    """Base for coil based entities."""
-
-    _attr_has_entity_name = True
-    _attr_entity_registry_enabled_default = False
-
-    def __init__(
-        self, coordinator: Coordinator, coil: Coil, entity_format: str
-    ) -> None:
-        """Initialize base entity."""
-        super().__init__(coordinator, {coil.address})
-        self.entity_id = async_generate_entity_id(
-            entity_format, coil.name, hass=coordinator.hass
-        )
-        self._attr_name = coil.title
-        self._attr_unique_id = f"{coordinator.unique_id}-{coil.address}"
-        self._attr_device_info = coordinator.device_info
-        self._coil = coil
-
-    @property
-    def available(self) -> bool:
-        """Return if entity is available."""
-        return self.coordinator.last_update_success and self._coil.address in (
-            self.coordinator.data or {}
-        )
-
-    def _async_read_coil(self, data: CoilData):
-        """Update state of entity based on coil data."""
-
-    async def _async_write_coil(self, value: float | str):
-        """Write coil and update state."""
-        await self.coordinator.async_write_coil(self._coil, value)
-
-    def _handle_coordinator_update(self) -> None:
-        data = self.coordinator.data.get(self._coil.address)
-        if data is not None:
-            self._async_read_coil(data)
-        self.async_write_ha_state()
diff --git a/homeassistant/components/nibe_heatpump/entity.py b/homeassistant/components/nibe_heatpump/entity.py
new file mode 100644
index 00000000000..3cbc8af32a3
--- /dev/null
+++ b/homeassistant/components/nibe_heatpump/entity.py
@@ -0,0 +1,50 @@
+"""The Nibe Heat Pump coordinator."""
+
+from __future__ import annotations
+
+from nibe.coil import Coil, CoilData
+
+from homeassistant.helpers.entity import async_generate_entity_id
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .coordinator import CoilCoordinator
+
+
+class CoilEntity(CoordinatorEntity[CoilCoordinator]):
+    """Base for coil based entities."""
+
+    _attr_has_entity_name = True
+    _attr_entity_registry_enabled_default = False
+
+    def __init__(
+        self, coordinator: CoilCoordinator, coil: Coil, entity_format: str
+    ) -> None:
+        """Initialize base entity."""
+        super().__init__(coordinator, {coil.address})
+        self.entity_id = async_generate_entity_id(
+            entity_format, coil.name, hass=coordinator.hass
+        )
+        self._attr_name = coil.title
+        self._attr_unique_id = f"{coordinator.unique_id}-{coil.address}"
+        self._attr_device_info = coordinator.device_info
+        self._coil = coil
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return self.coordinator.last_update_success and self._coil.address in (
+            self.coordinator.data or {}
+        )
+
+    def _async_read_coil(self, data: CoilData):
+        """Update state of entity based on coil data."""
+
+    async def _async_write_coil(self, value: float | str):
+        """Write coil and update state."""
+        await self.coordinator.async_write_coil(self._coil, value)
+
+    def _handle_coordinator_update(self) -> None:
+        data = self.coordinator.data.get(self._coil.address)
+        if data is not None:
+            self._async_read_coil(data)
+        self.async_write_ha_state()
diff --git a/homeassistant/components/nibe_heatpump/number.py b/homeassistant/components/nibe_heatpump/number.py
index 509f3364fee..cb379139eed 100644
--- a/homeassistant/components/nibe_heatpump/number.py
+++ b/homeassistant/components/nibe_heatpump/number.py
@@ -11,7 +11,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import CoilEntity, Coordinator
+from .coordinator import CoilCoordinator
+from .entity import CoilEntity
 
 
 async def async_setup_entry(
@@ -21,7 +22,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     async_add_entities(
         Number(coordinator, coil)
@@ -44,7 +45,7 @@ class Number(CoilEntity, NumberEntity):
 
     _attr_entity_category = EntityCategory.CONFIG
 
-    def __init__(self, coordinator: Coordinator, coil: Coil) -> None:
+    def __init__(self, coordinator: CoilCoordinator, coil: Coil) -> None:
         """Initialize entity."""
         super().__init__(coordinator, coil, ENTITY_ID_FORMAT)
         if coil.min is None or coil.max is None:
diff --git a/homeassistant/components/nibe_heatpump/select.py b/homeassistant/components/nibe_heatpump/select.py
index 07c958885b8..3aecff94649 100644
--- a/homeassistant/components/nibe_heatpump/select.py
+++ b/homeassistant/components/nibe_heatpump/select.py
@@ -11,7 +11,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import CoilEntity, Coordinator
+from .coordinator import CoilCoordinator
+from .entity import CoilEntity
 
 
 async def async_setup_entry(
@@ -21,7 +22,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     async_add_entities(
         Select(coordinator, coil)
@@ -35,7 +36,7 @@ class Select(CoilEntity, SelectEntity):
 
     _attr_entity_category = EntityCategory.CONFIG
 
-    def __init__(self, coordinator: Coordinator, coil: Coil) -> None:
+    def __init__(self, coordinator: CoilCoordinator, coil: Coil) -> None:
         """Initialize entity."""
         assert coil.mappings
         super().__init__(coordinator, coil, ENTITY_ID_FORMAT)
diff --git a/homeassistant/components/nibe_heatpump/sensor.py b/homeassistant/components/nibe_heatpump/sensor.py
index c6bac0323b9..d34fed50977 100644
--- a/homeassistant/components/nibe_heatpump/sensor.py
+++ b/homeassistant/components/nibe_heatpump/sensor.py
@@ -26,7 +26,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import CoilEntity, Coordinator
+from .coordinator import CoilCoordinator
+from .entity import CoilEntity
 
 UNIT_DESCRIPTIONS = {
     "°C": SensorEntityDescription(
@@ -130,7 +131,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     async_add_entities(
         Sensor(coordinator, coil, UNIT_DESCRIPTIONS.get(coil.unit))
@@ -144,7 +145,7 @@ class Sensor(CoilEntity, SensorEntity):
 
     def __init__(
         self,
-        coordinator: Coordinator,
+        coordinator: CoilCoordinator,
         coil: Coil,
         entity_description: SensorEntityDescription | None,
     ) -> None:
diff --git a/homeassistant/components/nibe_heatpump/switch.py b/homeassistant/components/nibe_heatpump/switch.py
index 594a8078b76..72b7c20c7b3 100644
--- a/homeassistant/components/nibe_heatpump/switch.py
+++ b/homeassistant/components/nibe_heatpump/switch.py
@@ -13,7 +13,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import CoilEntity, Coordinator
+from .coordinator import CoilCoordinator
+from .entity import CoilEntity
 
 
 async def async_setup_entry(
@@ -23,7 +24,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     async_add_entities(
         Switch(coordinator, coil)
@@ -37,7 +38,7 @@ class Switch(CoilEntity, SwitchEntity):
 
     _attr_entity_category = EntityCategory.CONFIG
 
-    def __init__(self, coordinator: Coordinator, coil: Coil) -> None:
+    def __init__(self, coordinator: CoilCoordinator, coil: Coil) -> None:
         """Initialize entity."""
         super().__init__(coordinator, coil, ENTITY_ID_FORMAT)
 
diff --git a/homeassistant/components/nibe_heatpump/water_heater.py b/homeassistant/components/nibe_heatpump/water_heater.py
index c60f5b6e3b2..f53df596d27 100644
--- a/homeassistant/components/nibe_heatpump/water_heater.py
+++ b/homeassistant/components/nibe_heatpump/water_heater.py
@@ -26,7 +26,7 @@ from .const import (
     VALUES_TEMPORARY_LUX_INACTIVE,
     VALUES_TEMPORARY_LUX_ONE_TIME_INCREASE,
 )
-from .coordinator import Coordinator
+from .coordinator import CoilCoordinator
 
 
 async def async_setup_entry(
@@ -36,7 +36,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up platform."""
 
-    coordinator: Coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: CoilCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     def water_heaters():
         for key, group in WATER_HEATER_COILGROUPS.get(coordinator.series, ()).items():
@@ -48,7 +48,7 @@ async def async_setup_entry(
     async_add_entities(water_heaters())
 
 
-class WaterHeater(CoordinatorEntity[Coordinator], WaterHeaterEntity):
+class WaterHeater(CoordinatorEntity[CoilCoordinator], WaterHeaterEntity):
     """Sensor entity."""
 
     _attr_entity_category = None
@@ -59,7 +59,7 @@ class WaterHeater(CoordinatorEntity[Coordinator], WaterHeaterEntity):
 
     def __init__(
         self,
-        coordinator: Coordinator,
+        coordinator: CoilCoordinator,
         key: str,
         desc: WaterHeaterCoilGroup,
     ) -> None: