diff --git a/homeassistant/components/iotty/__init__.py b/homeassistant/components/iotty/__init__.py index 804f3f40196..c9eb2639348 100644 --- a/homeassistant/components/iotty/__init__.py +++ b/homeassistant/components/iotty/__init__.py @@ -2,12 +2,8 @@ from __future__ import annotations -from dataclasses import dataclass import logging -from iottycloud.device import Device - -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.helpers.config_entry_oauth2_flow import ( @@ -15,22 +11,16 @@ from homeassistant.helpers.config_entry_oauth2_flow import ( async_get_config_entry_implementation, ) -from . import coordinator +from .coordinator import ( + IottyConfigEntry, + IottyConfigEntryData, + IottyDataUpdateCoordinator, +) _LOGGER = logging.getLogger(__name__) PLATFORMS: list[Platform] = [Platform.COVER, Platform.SWITCH] -type IottyConfigEntry = ConfigEntry[IottyConfigEntryData] - - -@dataclass -class IottyConfigEntryData: - """Contains config entry data for iotty.""" - - known_devices: set[Device] - coordinator: coordinator.IottyDataUpdateCoordinator - async def async_setup_entry(hass: HomeAssistant, entry: IottyConfigEntry) -> bool: """Set up iotty from a config entry.""" @@ -39,9 +29,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: IottyConfigEntry) -> boo implementation = await async_get_config_entry_implementation(hass, entry) session = OAuth2Session(hass, entry, implementation) - data_update_coordinator = coordinator.IottyDataUpdateCoordinator( - hass, entry, session - ) + data_update_coordinator = IottyDataUpdateCoordinator(hass, entry, session) entry.runtime_data = IottyConfigEntryData(set(), data_update_coordinator) @@ -51,6 +39,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: IottyConfigEntry) -> boo return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: IottyConfigEntry) -> bool: """Unload a config entry.""" return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/iotty/coordinator.py b/homeassistant/components/iotty/coordinator.py index 420248f7724..af870c347bd 100644 --- a/homeassistant/components/iotty/coordinator.py +++ b/homeassistant/components/iotty/coordinator.py @@ -32,16 +32,27 @@ class IottyData: devices: list[Device] +@dataclass +class IottyConfigEntryData: + """Contains config entry data for iotty.""" + + known_devices: set[Device] + coordinator: IottyDataUpdateCoordinator + + +type IottyConfigEntry = ConfigEntry[IottyConfigEntryData] + + class IottyDataUpdateCoordinator(DataUpdateCoordinator[IottyData]): """Class to manage fetching Iotty data.""" - config_entry: ConfigEntry + config_entry: IottyConfigEntry _entities: dict[str, Entity] _devices: list[Device] _device_registry: dr.DeviceRegistry def __init__( - self, hass: HomeAssistant, entry: ConfigEntry, session: OAuth2Session + self, hass: HomeAssistant, entry: IottyConfigEntry, session: OAuth2Session ) -> None: """Initialize the coordinator.""" _LOGGER.debug("Initializing iotty data update coordinator") @@ -49,11 +60,11 @@ class IottyDataUpdateCoordinator(DataUpdateCoordinator[IottyData]): super().__init__( hass, _LOGGER, + config_entry=entry, name=f"{DOMAIN}_coordinator", update_interval=UPDATE_INTERVAL, ) - self.config_entry = entry self._entities = {} self._devices = [] self.iotty = api.IottyProxy( diff --git a/homeassistant/components/iotty/cover.py b/homeassistant/components/iotty/cover.py index 50a4a1deeba..31d363868db 100644 --- a/homeassistant/components/iotty/cover.py +++ b/homeassistant/components/iotty/cover.py @@ -18,9 +18,8 @@ from homeassistant.components.cover import ( from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import IottyConfigEntry from .api import IottyProxy -from .coordinator import IottyDataUpdateCoordinator +from .coordinator import IottyConfigEntry, IottyDataUpdateCoordinator from .entity import IottyEntity _LOGGER = logging.getLogger(__name__) diff --git a/homeassistant/components/iotty/switch.py b/homeassistant/components/iotty/switch.py index b06e3ea308d..a748ac10783 100644 --- a/homeassistant/components/iotty/switch.py +++ b/homeassistant/components/iotty/switch.py @@ -22,9 +22,8 @@ from homeassistant.components.switch import ( from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import IottyConfigEntry from .api import IottyProxy -from .coordinator import IottyDataUpdateCoordinator +from .coordinator import IottyConfigEntry, IottyDataUpdateCoordinator from .entity import IottyEntity _LOGGER = logging.getLogger(__name__) diff --git a/tests/components/iotty/conftest.py b/tests/components/iotty/conftest.py index 51a23bf18c7..1ce645b402e 100644 --- a/tests/components/iotty/conftest.py +++ b/tests/components/iotty/conftest.py @@ -169,7 +169,7 @@ def mock_iotty() -> Generator[MagicMock]: def mock_coordinator() -> Generator[MagicMock]: """Mock IottyDataUpdateCoordinator.""" with patch( - "homeassistant.components.iotty.coordinator.IottyDataUpdateCoordinator", + "homeassistant.components.iotty.IottyDataUpdateCoordinator", autospec=True, ) as coordinator_mock: yield coordinator_mock