Move p1_monitor coordinator to separate module (home-assistant#117562)

cereal2nd · May 17, 2024 · fce4263 · fce4263
1 parent c419624
commit fce4263
Show file tree

Hide file tree

Showing 7 changed files with 93 additions and 81 deletions.
diff --git a/homeassistant/components/p1_monitor/__init__.py b/homeassistant/components/p1_monitor/__init__.py
@@ -2,34 +2,13 @@
 
 from __future__ import annotations
 
-from typing import TypedDict
-
-from p1monitor import (
-    P1Monitor,
-    P1MonitorConnectionError,
-    P1MonitorNoDataError,
-    Phases,
-    Settings,
-    SmartMeter,
-    WaterMeter,
-)
-
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_HOST, Platform
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import (
-    DOMAIN,
-    LOGGER,
-    SCAN_INTERVAL,
-    SERVICE_PHASES,
-    SERVICE_SETTINGS,
-    SERVICE_SMARTMETER,
-    SERVICE_WATERMETER,
-)
+from .const import DOMAIN
+from .coordinator import P1MonitorDataUpdateCoordinator
 
 PLATFORMS = [Platform.SENSOR]
 
@@ -57,55 +36,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if unload_ok:
         del hass.data[DOMAIN][entry.entry_id]
     return unload_ok
-
-
-class P1MonitorData(TypedDict):
-    """Class for defining data in dict."""
-
-    smartmeter: SmartMeter
-    phases: Phases
-    settings: Settings
-    watermeter: WaterMeter | None
-
-
-class P1MonitorDataUpdateCoordinator(DataUpdateCoordinator[P1MonitorData]):  # pylint: disable=hass-enforce-coordinator-module
-    """Class to manage fetching P1 Monitor data from single endpoint."""
-
-    config_entry: ConfigEntry
-    has_water_meter: bool | None = None
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-    ) -> None:
-        """Initialize global P1 Monitor data updater."""
-        super().__init__(
-            hass,
-            LOGGER,
-            name=DOMAIN,
-            update_interval=SCAN_INTERVAL,
-        )
-
-        self.p1monitor = P1Monitor(
-            self.config_entry.data[CONF_HOST], session=async_get_clientsession(hass)
-        )
-
-    async def _async_update_data(self) -> P1MonitorData:
-        """Fetch data from P1 Monitor."""
-        data: P1MonitorData = {
-            SERVICE_SMARTMETER: await self.p1monitor.smartmeter(),
-            SERVICE_PHASES: await self.p1monitor.phases(),
-            SERVICE_SETTINGS: await self.p1monitor.settings(),
-            SERVICE_WATERMETER: None,
-        }
-
-        if self.has_water_meter or self.has_water_meter is None:
-            try:
-                data[SERVICE_WATERMETER] = await self.p1monitor.watermeter()
-                self.has_water_meter = True
-            except (P1MonitorNoDataError, P1MonitorConnectionError):
-                LOGGER.debug("No water meter data received from P1 Monitor")
-                if self.has_water_meter is None:
-                    self.has_water_meter = False
-
-        return data
diff --git a/homeassistant/components/p1_monitor/coordinator.py b/homeassistant/components/p1_monitor/coordinator.py
@@ -0,0 +1,83 @@
+"""Coordinator for the P1 Monitor integration."""
+
+from __future__ import annotations
+
+from typing import TypedDict
+
+from p1monitor import (
+    P1Monitor,
+    P1MonitorConnectionError,
+    P1MonitorNoDataError,
+    Phases,
+    Settings,
+    SmartMeter,
+    WaterMeter,
+)
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_HOST
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+
+from .const import (
+    DOMAIN,
+    LOGGER,
+    SCAN_INTERVAL,
+    SERVICE_PHASES,
+    SERVICE_SETTINGS,
+    SERVICE_SMARTMETER,
+    SERVICE_WATERMETER,
+)
+
+
+class P1MonitorData(TypedDict):
+    """Class for defining data in dict."""
+
+    smartmeter: SmartMeter
+    phases: Phases
+    settings: Settings
+    watermeter: WaterMeter | None
+
+
+class P1MonitorDataUpdateCoordinator(DataUpdateCoordinator[P1MonitorData]):
+    """Class to manage fetching P1 Monitor data from single endpoint."""
+
+    config_entry: ConfigEntry
+    has_water_meter: bool | None = None
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+    ) -> None:
+        """Initialize global P1 Monitor data updater."""
+        super().__init__(
+            hass,
+            LOGGER,
+            name=DOMAIN,
+            update_interval=SCAN_INTERVAL,
+        )
+
+        self.p1monitor = P1Monitor(
+            self.config_entry.data[CONF_HOST], session=async_get_clientsession(hass)
+        )
+
+    async def _async_update_data(self) -> P1MonitorData:
+        """Fetch data from P1 Monitor."""
+        data: P1MonitorData = {
+            SERVICE_SMARTMETER: await self.p1monitor.smartmeter(),
+            SERVICE_PHASES: await self.p1monitor.phases(),
+            SERVICE_SETTINGS: await self.p1monitor.settings(),
+            SERVICE_WATERMETER: None,
+        }
+
+        if self.has_water_meter or self.has_water_meter is None:
+            try:
+                data[SERVICE_WATERMETER] = await self.p1monitor.watermeter()
+                self.has_water_meter = True
+            except (P1MonitorNoDataError, P1MonitorConnectionError):
+                LOGGER.debug("No water meter data received from P1 Monitor")
+                if self.has_water_meter is None:
+                    self.has_water_meter = False
+
+        return data
diff --git a/homeassistant/components/p1_monitor/diagnostics.py b/homeassistant/components/p1_monitor/diagnostics.py
@@ -10,14 +10,14 @@
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
 
-from . import P1MonitorDataUpdateCoordinator
 from .const import (
     DOMAIN,
     SERVICE_PHASES,
     SERVICE_SETTINGS,
     SERVICE_SMARTMETER,
     SERVICE_WATERMETER,
 )
+from .coordinator import P1MonitorDataUpdateCoordinator
 
 if TYPE_CHECKING:
     from _typeshed import DataclassInstance

diff --git a/homeassistant/components/p1_monitor/sensor.py b/homeassistant/components/p1_monitor/sensor.py
@@ -26,14 +26,14 @@
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import P1MonitorDataUpdateCoordinator
 from .const import (
     DOMAIN,
     SERVICE_PHASES,
     SERVICE_SETTINGS,
     SERVICE_SMARTMETER,
     SERVICE_WATERMETER,
 )
+from .coordinator import P1MonitorDataUpdateCoordinator
 
 SENSORS_SMARTMETER: tuple[SensorEntityDescription, ...] = (
     SensorEntityDescription(

diff --git a/tests/components/p1_monitor/conftest.py b/tests/components/p1_monitor/conftest.py
@@ -27,7 +27,9 @@ def mock_config_entry() -> MockConfigEntry:
 @pytest.fixture
 def mock_p1monitor():
     """Return a mocked P1 Monitor client."""
-    with patch("homeassistant.components.p1_monitor.P1Monitor") as p1monitor_mock:
+    with patch(
+        "homeassistant.components.p1_monitor.coordinator.P1Monitor"
+    ) as p1monitor_mock:
         client = p1monitor_mock.return_value
         client.smartmeter = AsyncMock(
             return_value=SmartMeter.from_dict(

diff --git a/tests/components/p1_monitor/test_config_flow.py b/tests/components/p1_monitor/test_config_flow.py
@@ -44,7 +44,7 @@ async def test_full_user_flow(hass: HomeAssistant) -> None:
 async def test_api_error(hass: HomeAssistant) -> None:
     """Test we handle cannot connect error."""
     with patch(
-        "homeassistant.components.p1_monitor.P1Monitor.smartmeter",
+        "homeassistant.components.p1_monitor.coordinator.P1Monitor.smartmeter",
         side_effect=P1MonitorError,
     ):
         result = await hass.config_entries.flow.async_init(

diff --git a/tests/components/p1_monitor/test_init.py b/tests/components/p1_monitor/test_init.py
@@ -29,7 +29,7 @@ async def test_load_unload_config_entry(
 
 
 @patch(
-    "homeassistant.components.p1_monitor.P1Monitor._request",
+    "homeassistant.components.p1_monitor.coordinator.P1Monitor._request",
     side_effect=P1MonitorConnectionError,
 )
 async def test_config_entry_not_ready(