Move brunt coordinator to separate module (#129090)

homeassistant/components/brunt/__init__.py

@@ -2,79 +2,22 @@
 
 from __future__ import annotations
 
-from asyncio import timeout
-import logging
-
-from aiohttp.client_exceptions import ClientResponseError, ServerDisconnectedError
-from brunt import BruntClientAsync, Thing
-
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DATA_BAPI, DATA_COOR, DOMAIN, PLATFORMS, REGULAR_INTERVAL
-
-_LOGGER = logging.getLogger(__name__)
+from .const import PLATFORMS
+from .coordinator import BruntConfigEntry, BruntCoordinator
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: BruntConfigEntry) -> bool:
     """Set up Brunt using config flow."""
-    session = async_get_clientsession(hass)
-    bapi = BruntClientAsync(
-        username=entry.data[CONF_USERNAME],
-        password=entry.data[CONF_PASSWORD],
-        session=session,
-    )
-    try:
-        await bapi.async_login()
-    except ServerDisconnectedError as exc:
-        raise ConfigEntryNotReady("Brunt not ready to connect.") from exc
-    except ClientResponseError as exc:
-        raise ConfigEntryAuthFailed(
-            f"Brunt could not connect with username: {entry.data[CONF_USERNAME]}."
-        ) from exc
-
-    async def async_update_data() -> dict[str | None, Thing]:
-        """Fetch data from the Brunt endpoint for all Things.
-
-        Error 403 is the API response for any kind of authentication error (failed password or email)
-        Error 401 is the API response for things that are not part of the account, could happen when a device is deleted from the account.
-        """
-        try:
-            async with timeout(10):
-                things = await bapi.async_get_things(force=True)
-                return {thing.serial: thing for thing in things}
-        except ServerDisconnectedError as err:
-            raise UpdateFailed(f"Error communicating with API: {err}") from err
-        except ClientResponseError as err:
-            if err.status == 403:
-                raise ConfigEntryAuthFailed from err
-            if err.status == 401:
-                _LOGGER.warning("Device not found, will reload Brunt integration")
-                await hass.config_entries.async_reload(entry.entry_id)
-            raise UpdateFailed from err
-
-    coordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        name="brunt",
-        update_method=async_update_data,
-        update_interval=REGULAR_INTERVAL,
-    )
+    coordinator = BruntCoordinator(hass, entry)
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {DATA_BAPI: bapi, DATA_COOR: coordinator}
+    entry.runtime_data = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: BruntConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/brunt/const.py

@@ -10,8 +10,6 @@ NOTIFICATION_ID = "brunt_notification"
 NOTIFICATION_TITLE = "Brunt Cover Setup"
 ATTRIBUTION = "Based on an unofficial Brunt SDK."
 PLATFORMS = [Platform.COVER]
-DATA_BAPI = "bapi"
-DATA_COOR = "coordinator"
 
 CLOSED_POSITION = 0
 OPEN_POSITION = 100

homeassistant/components/brunt/coordinator.py

@@ -0,0 +1,80 @@
+"""The brunt component."""
+
+from __future__ import annotations
+
+from asyncio import timeout
+import logging
+
+from aiohttp.client_exceptions import ClientResponseError, ServerDisconnectedError
+from brunt import BruntClientAsync, Thing
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import REGULAR_INTERVAL
+
+_LOGGER = logging.getLogger(__name__)
+
+type BruntConfigEntry = ConfigEntry[BruntCoordinator]
+
+
+class BruntCoordinator(DataUpdateCoordinator[dict[str | None, Thing]]):
+    """Config entry data."""
+
+    bapi: BruntClientAsync
+    config_entry: BruntConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: BruntConfigEntry,
+    ) -> None:
+        """Initialize the Brunt coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=config_entry,
+            name="brunt",
+            update_interval=REGULAR_INTERVAL,
+        )
+
+    async def _async_setup(self) -> None:
+        session = async_get_clientsession(self.hass)
+
+        self.bapi = BruntClientAsync(
+            username=self.config_entry.data[CONF_USERNAME],
+            password=self.config_entry.data[CONF_PASSWORD],
+            session=session,
+        )
+        try:
+            await self.bapi.async_login()
+        except ServerDisconnectedError as exc:
+            raise ConfigEntryNotReady("Brunt not ready to connect.") from exc
+        except ClientResponseError as exc:
+            raise ConfigEntryAuthFailed(
+                f"Brunt could not connect with username: {self.config_entry.data[CONF_USERNAME]}."
+            ) from exc
+
+    async def _async_update_data(self) -> dict[str | None, Thing]:
+        """Fetch data from the Brunt endpoint for all Things.
+
+        Error 403 is the API response for any kind of authentication error (failed password or email)
+        Error 401 is the API response for things that are not part of the account, could happen when a device is deleted from the account.
+        """
+        try:
+            async with timeout(10):
+                things = await self.bapi.async_get_things(force=True)
+                return {thing.serial: thing for thing in things}
+        except ServerDisconnectedError as err:
+            raise UpdateFailed(f"Error communicating with API: {err}") from err
+        except ClientResponseError as err:
+            if err.status == 403:
+                raise ConfigEntryAuthFailed from err
+            if err.status == 401:
+                _LOGGER.warning("Device not found, will reload Brunt integration")
+                await self.hass.config_entries.async_reload(self.config_entry.entry_id)
+            raise UpdateFailed from err

homeassistant/components/brunt/cover.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from typing import Any
 
 from aiohttp.client_exceptions import ClientResponseError
-from brunt import BruntClientAsync, Thing
+from brunt import Thing
 
 from homeassistant.components.cover import (
     ATTR_POSITION,
@@ -13,49 +13,39 @@ from homeassistant.components.cover import (
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import (
     ATTR_REQUEST_POSITION,
     ATTRIBUTION,
     CLOSED_POSITION,
-    DATA_BAPI,
-    DATA_COOR,
     DOMAIN,
     FAST_INTERVAL,
     OPEN_POSITION,
     REGULAR_INTERVAL,
 )
+from .coordinator import BruntConfigEntry, BruntCoordinator
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: BruntConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the brunt platform."""
-    bapi: BruntClientAsync = hass.data[DOMAIN][entry.entry_id][DATA_BAPI]
-    coordinator: DataUpdateCoordinator[dict[str | None, Thing]] = hass.data[DOMAIN][
-        entry.entry_id
-    ][DATA_COOR]
+    coordinator = entry.runtime_data
 
     async_add_entities(
-        BruntDevice(coordinator, serial, thing, bapi, entry.entry_id)
+        BruntDevice(coordinator, serial, thing, entry.entry_id)
         for serial, thing in coordinator.data.items()
     )
 
 
-class BruntDevice(
-    CoordinatorEntity[DataUpdateCoordinator[dict[str | None, Thing]]], CoverEntity
-):
+class BruntDevice(CoordinatorEntity[BruntCoordinator], CoverEntity):
     """Representation of a Brunt cover device.
 
     Contains the common logic for all Brunt devices.
@@ -73,16 +63,14 @@ class BruntDevice(
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[dict[str | None, Thing]],
+        coordinator: BruntCoordinator,
         serial: str | None,
         thing: Thing,
-        bapi: BruntClientAsync,
         entry_id: str,
     ) -> None:
         """Init the Brunt device."""
         super().__init__(coordinator)
         self._attr_unique_id = serial
-        self._bapi = bapi
         self._thing = thing
         self._entry_id = entry_id
 
@@ -167,7 +155,7 @@ class BruntDevice(
     async def _async_update_cover(self, position: int) -> None:
         """Set the cover to the new position and wait for the update to be reflected."""
         try:
-            await self._bapi.async_change_request_position(
+            await self.coordinator.bapi.async_change_request_position(
                 position, thing_uri=self._thing.thing_uri
             )
         except ClientResponseError as exc:
@@ -182,7 +170,7 @@ class BruntDevice(
         """Update the update interval after each refresh."""
         if (
             self.request_cover_position
-            == self._bapi.last_requested_positions[self._thing.thing_uri]
+            == self.coordinator.bapi.last_requested_positions[self._thing.thing_uri]
             and self.move_state == 0
         ):
             self.coordinator.update_interval = REGULAR_INTERVAL