Move omnilogic coordinator to separate module (#118014)

home-assistant · May 24, 2024 · 8da799e · 8da799e
1 parent 301c484
commit 8da799e
Show file tree

Hide file tree

Showing 7 changed files with 77 additions and 71 deletions.
diff --git a/.coveragerc b/.coveragerc
@@ -932,7 +932,7 @@ omit =
     homeassistant/components/ohmconnect/sensor.py
     homeassistant/components/ombi/*
     homeassistant/components/omnilogic/__init__.py
-    homeassistant/components/omnilogic/common.py
+    homeassistant/components/omnilogic/coordinator.py
     homeassistant/components/omnilogic/sensor.py
     homeassistant/components/omnilogic/switch.py
     homeassistant/components/ondilo_ico/__init__.py

diff --git a/homeassistant/components/omnilogic/__init__.py b/homeassistant/components/omnilogic/__init__.py
@@ -10,14 +10,14 @@
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client
 
-from .common import OmniLogicUpdateCoordinator
 from .const import (
     CONF_SCAN_INTERVAL,
     COORDINATOR,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
     OMNI_API,
 )
+from .coordinator import OmniLogicUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 

diff --git a/homeassistant/components/omnilogic/common.py b/homeassistant/components/omnilogic/common.py
@@ -1,75 +1,12 @@
 """Common classes and elements for Omnilogic Integration."""
 
-from datetime import timedelta
-import logging
 from typing import Any
 
-from omnilogic import OmniLogic, OmniLogicException
-
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-    UpdateFailed,
-)
-
-from .const import ALL_ITEM_KINDS, DOMAIN
-
-_LOGGER = logging.getLogger(__name__)
-
-
-class OmniLogicUpdateCoordinator(DataUpdateCoordinator[dict[tuple, dict[str, Any]]]):  # pylint: disable=hass-enforce-coordinator-module
-    """Class to manage fetching update data from single endpoint."""
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        api: OmniLogic,
-        name: str,
-        config_entry: ConfigEntry,
-        polling_interval: int,
-    ) -> None:
-        """Initialize the global Omnilogic data updater."""
-        self.api = api
-        self.config_entry = config_entry
-
-        super().__init__(
-            hass=hass,
-            logger=_LOGGER,
-            name=name,
-            update_interval=timedelta(seconds=polling_interval),
-        )
-
-    async def _async_update_data(self):
-        """Fetch data from OmniLogic."""
-        try:
-            data = await self.api.get_telemetry_data()
-
-        except OmniLogicException as error:
-            raise UpdateFailed(f"Error updating from OmniLogic: {error}") from error
-
-        parsed_data = {}
-
-        def get_item_data(item, item_kind, current_id, data):
-            """Get data per kind of Omnilogic API item."""
-            if isinstance(item, list):
-                for single_item in item:
-                    data = get_item_data(single_item, item_kind, current_id, data)
-
-            if "systemId" in item:
-                system_id = item["systemId"]
-                current_id = (*current_id, item_kind, system_id)
-                data[current_id] = item
-
-            for kind in ALL_ITEM_KINDS:
-                if kind in item:
-                    data = get_item_data(item[kind], kind, current_id, data)
-
-            return data
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-        return get_item_data(data, "Backyard", (), parsed_data)
+from .const import DOMAIN
+from .coordinator import OmniLogicUpdateCoordinator
 
 
 class OmniLogicEntity(CoordinatorEntity[OmniLogicUpdateCoordinator]):

diff --git a/homeassistant/components/omnilogic/coordinator.py b/homeassistant/components/omnilogic/coordinator.py
@@ -0,0 +1,67 @@
+"""Coordinator for the Omnilogic Integration."""
+
+from datetime import timedelta
+import logging
+from typing import Any
+
+from omnilogic import OmniLogic, OmniLogicException
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import ALL_ITEM_KINDS
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class OmniLogicUpdateCoordinator(DataUpdateCoordinator[dict[tuple, dict[str, Any]]]):
+    """Class to manage fetching update data from single endpoint."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        api: OmniLogic,
+        name: str,
+        config_entry: ConfigEntry,
+        polling_interval: int,
+    ) -> None:
+        """Initialize the global Omnilogic data updater."""
+        self.api = api
+        self.config_entry = config_entry
+
+        super().__init__(
+            hass=hass,
+            logger=_LOGGER,
+            name=name,
+            update_interval=timedelta(seconds=polling_interval),
+        )
+
+    async def _async_update_data(self):
+        """Fetch data from OmniLogic."""
+        try:
+            data = await self.api.get_telemetry_data()
+
+        except OmniLogicException as error:
+            raise UpdateFailed(f"Error updating from OmniLogic: {error}") from error
+
+        parsed_data = {}
+
+        def get_item_data(item, item_kind, current_id, data):
+            """Get data per kind of Omnilogic API item."""
+            if isinstance(item, list):
+                for single_item in item:
+                    data = get_item_data(single_item, item_kind, current_id, data)
+
+            if "systemId" in item:
+                system_id = item["systemId"]
+                current_id = (*current_id, item_kind, system_id)
+                data[current_id] = item
+
+            for kind in ALL_ITEM_KINDS:
+                if kind in item:
+                    data = get_item_data(item[kind], kind, current_id, data)
+
+            return data
+
+        return get_item_data(data, "Backyard", (), parsed_data)
diff --git a/homeassistant/components/omnilogic/sensor.py b/homeassistant/components/omnilogic/sensor.py
@@ -15,8 +15,9 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .common import OmniLogicEntity, OmniLogicUpdateCoordinator, check_guard
+from .common import OmniLogicEntity, check_guard
 from .const import COORDINATOR, DEFAULT_PH_OFFSET, DOMAIN, PUMP_TYPES
+from .coordinator import OmniLogicUpdateCoordinator
 
 
 async def async_setup_entry(

diff --git a/homeassistant/components/omnilogic/switch.py b/homeassistant/components/omnilogic/switch.py
@@ -12,8 +12,9 @@
 from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .common import OmniLogicEntity, OmniLogicUpdateCoordinator, check_guard
+from .common import OmniLogicEntity, check_guard
 from .const import COORDINATOR, DOMAIN, PUMP_TYPES
+from .coordinator import OmniLogicUpdateCoordinator
 
 SERVICE_SET_SPEED = "set_pump_speed"
 OMNILOGIC_SWITCH_OFF = 7

diff --git a/tests/components/omnilogic/__init__.py b/tests/components/omnilogic/__init__.py
@@ -23,7 +23,7 @@ async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
             return_value={},
         ),
         patch(
-            "homeassistant.components.omnilogic.common.OmniLogicUpdateCoordinator._async_update_data",
+            "homeassistant.components.omnilogic.coordinator.OmniLogicUpdateCoordinator._async_update_data",
             return_value=TELEMETRY,
         ),
     ):