diff --git a/homeassistant/components/brother/__init__.py b/homeassistant/components/brother/__init__.py index 0bd49ed5d7a..08376574dcf 100644 --- a/homeassistant/components/brother/__init__.py +++ b/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 diff --git a/homeassistant/components/brother/const.py b/homeassistant/components/brother/const.py index fda815ceee5..f8d29363acd 100644 --- a/homeassistant/components/brother/const.py +++ b/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) diff --git a/homeassistant/components/brother/coordinator.py b/homeassistant/components/brother/coordinator.py new file mode 100644 index 00000000000..69463d107e4 --- /dev/null +++ b/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 diff --git a/homeassistant/components/brother/diagnostics.py b/homeassistant/components/brother/diagnostics.py index ee5eedd84cb..d4a6c6c5400 100644 --- a/homeassistant/components/brother/diagnostics.py +++ b/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), diff --git a/homeassistant/components/brother/sensor.py b/homeassistant/components/brother/sensor.py index 6f56eb680be..e86eb59d6bc 100644 --- a/homeassistant/components/brother/sensor.py +++ b/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)