Move Brother DataUpdateCoordinator to the coordinator module (#116772)

* Store data in config_entry.runtime_data

* Update homeassistant/components/brother/__init__.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* Fix setdefault

* Do not include ignored and disabled

* Check loaded entries

---------

Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>
Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/brother/__init__.py

@@ -2,29 +2,23 @@
 
 from __future__ import annotations
 
-from asyncio import timeout
-from datetime import timedelta
-import logging
+from brother import Brother, SnmpError
 
-from brother import Brother, BrotherSensors, SnmpError, UnsupportedModelError
-
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import CONF_HOST, CONF_TYPE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DATA_CONFIG_ENTRY, DOMAIN, SNMP
+from .const import DOMAIN, SNMP
+from .coordinator import BrotherDataUpdateCoordinator
 from .utils import get_snmp_engine
 
 PLATFORMS = [Platform.SENSOR]
 
-SCAN_INTERVAL = timedelta(seconds=30)
-
-_LOGGER = logging.getLogger(__name__)
+BrotherConfigEntry = ConfigEntry[BrotherDataUpdateCoordinator]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: BrotherConfigEntry) -> bool:
     """Set up Brother from a config entry."""
     host = entry.data[CONF_HOST]
     printer_type = entry.data[CONF_TYPE]
@@ -40,48 +34,25 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     coordinator = BrotherDataUpdateCoordinator(hass, brother)
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN].setdefault(DATA_CONFIG_ENTRY, {})
-    hass.data[DOMAIN][DATA_CONFIG_ENTRY][entry.entry_id] = coordinator
-    hass.data[DOMAIN][SNMP] = snmp_engine
+    entry.runtime_data = coordinator
+    hass.data.setdefault(DOMAIN, {SNMP: snmp_engine})
 
     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: BrotherConfigEntry) -> bool:
     """Unload a config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
-    if unload_ok:
-        hass.data[DOMAIN][DATA_CONFIG_ENTRY].pop(entry.entry_id)
-        if not hass.data[DOMAIN][DATA_CONFIG_ENTRY]:
-            hass.data[DOMAIN].pop(SNMP)
-            hass.data[DOMAIN].pop(DATA_CONFIG_ENTRY)
+    loaded_entries = [
+        entry
+        for entry in hass.config_entries.async_entries(DOMAIN)
+        if entry.state == ConfigEntryState.LOADED
+    ]
+    # We only want to remove the SNMP engine when unloading the last config entry
+    if unload_ok and len(loaded_entries) == 1:
+        hass.data[DOMAIN].pop(SNMP)
 
     return unload_ok
-
-
-class BrotherDataUpdateCoordinator(DataUpdateCoordinator[BrotherSensors]):  # pylint: disable=hass-enforce-coordinator-module
-    """Class to manage fetching Brother data from the printer."""
-
-    def __init__(self, hass: HomeAssistant, brother: Brother) -> None:
-        """Initialize."""
-        self.brother = brother
-
-        super().__init__(
-            hass,
-            _LOGGER,
-            name=DOMAIN,
-            update_interval=SCAN_INTERVAL,
-        )
-
-    async def _async_update_data(self) -> BrotherSensors:
-        """Update data via library."""
-        try:
-            async with timeout(20):
-                data = await self.brother.async_update()
-        except (ConnectionError, SnmpError, UnsupportedModelError) as error:
-            raise UpdateFailed(error) from error
-        return data

homeassistant/components/brother/const.py

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
+from datetime import timedelta
 from typing import Final
 
-DATA_CONFIG_ENTRY: Final = "config_entry"
-
 DOMAIN: Final = "brother"
 
 PRINTER_TYPES: Final = ["laser", "ink"]
 
 SNMP: Final = "snmp"
+
+UPDATE_INTERVAL = timedelta(seconds=30)

homeassistant/components/brother/coordinator.py

@@ -0,0 +1,37 @@
+"""Coordinator for Brother integration."""
+
+from asyncio import timeout
+import logging
+
+from brother import Brother, BrotherSensors, SnmpError, UnsupportedModelError
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN, UPDATE_INTERVAL
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class BrotherDataUpdateCoordinator(DataUpdateCoordinator[BrotherSensors]):
+    """Class to manage fetching Brother data from the printer."""
+
+    def __init__(self, hass: HomeAssistant, brother: Brother) -> None:
+        """Initialize."""
+        self.brother = brother
+
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=UPDATE_INTERVAL,
+        )
+
+    async def _async_update_data(self) -> BrotherSensors:
+        """Update data via library."""
+        try:
+            async with timeout(20):
+                data = await self.brother.async_update()
+        except (ConnectionError, SnmpError, UnsupportedModelError) as error:
+            raise UpdateFailed(error) from error
+        return data

homeassistant/components/brother/diagnostics.py

@@ -5,20 +5,16 @@ from __future__ import annotations
 from dataclasses import asdict
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from . import BrotherDataUpdateCoordinator
-from .const import DATA_CONFIG_ENTRY, DOMAIN
+from . import BrotherConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: BrotherConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: BrotherDataUpdateCoordinator = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
-        config_entry.entry_id
-    ]
+    coordinator = config_entry.runtime_data
 
     return {
         "info": dict(config_entry.data),

homeassistant/components/brother/sensor.py

@@ -16,7 +16,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
@@ -25,8 +24,8 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import BrotherDataUpdateCoordinator
-from .const import DATA_CONFIG_ENTRY, DOMAIN
+from . import BrotherConfigEntry, BrotherDataUpdateCoordinator
+from .const import DOMAIN
 
 ATTR_COUNTER = "counter"
 ATTR_REMAINING_PAGES = "remaining_pages"
@@ -318,11 +317,12 @@ SENSOR_TYPES: tuple[BrotherSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: BrotherConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Add Brother entities from a config_entry."""
-    coordinator = hass.data[DOMAIN][DATA_CONFIG_ENTRY][entry.entry_id]
-
+    coordinator = entry.runtime_data
     # Due to the change of the attribute name of one sensor, it is necessary to migrate
     # the unique_id to the new one.
     entity_registry = er.async_get(hass)