From 030fe6df4a422161355f731ae5998a4c954da603 Mon Sep 17 00:00:00 2001 From: Rami Mosleh Date: Thu, 13 Jun 2024 12:53:32 +0300 Subject: [PATCH] Store Mikrotik coordinator in runtime_data (#119594) --- homeassistant/components/mikrotik/__init__.py | 15 +++++++-------- .../components/mikrotik/device_tracker.py | 9 +++------ tests/components/mikrotik/test_init.py | 2 -- 3 files changed, 10 insertions(+), 16 deletions(-) diff --git a/homeassistant/components/mikrotik/__init__.py b/homeassistant/components/mikrotik/__init__.py index 8e5911677af..9f2b40bf1c8 100644 --- a/homeassistant/components/mikrotik/__init__.py +++ b/homeassistant/components/mikrotik/__init__.py @@ -14,8 +14,12 @@ CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False) PLATFORMS = [Platform.DEVICE_TRACKER] +type MikrotikConfigEntry = ConfigEntry[MikrotikDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: + +async def async_setup_entry( + hass: HomeAssistant, config_entry: MikrotikConfigEntry +) -> bool: """Set up the Mikrotik component.""" try: api = await hass.async_add_executor_job(get_api, dict(config_entry.data)) @@ -28,7 +32,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b await hass.async_add_executor_job(coordinator.api.get_hub_details) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = coordinator + config_entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS) @@ -47,9 +51,4 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms( - config_entry, PLATFORMS - ): - hass.data[DOMAIN].pop(config_entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS) diff --git a/homeassistant/components/mikrotik/device_tracker.py b/homeassistant/components/mikrotik/device_tracker.py index 073db547b4c..aa19da01369 100644 --- a/homeassistant/components/mikrotik/device_tracker.py +++ b/homeassistant/components/mikrotik/device_tracker.py @@ -9,26 +9,23 @@ from homeassistant.components.device_tracker import ( ScannerEntity, SourceType, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity import homeassistant.util.dt as dt_util -from .const import DOMAIN +from . import MikrotikConfigEntry from .coordinator import Device, MikrotikDataUpdateCoordinator async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: MikrotikConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up device tracker for Mikrotik component.""" - coordinator: MikrotikDataUpdateCoordinator = hass.data[DOMAIN][ - config_entry.entry_id - ] + coordinator = config_entry.runtime_data tracked: dict[str, MikrotikDataUpdateCoordinatorTracker] = {} diff --git a/tests/components/mikrotik/test_init.py b/tests/components/mikrotik/test_init.py index cc6a737e75a..97245480300 100644 --- a/tests/components/mikrotik/test_init.py +++ b/tests/components/mikrotik/test_init.py @@ -6,7 +6,6 @@ from librouteros.exceptions import ConnectionClosed, LibRouterosError import pytest from homeassistant.components import mikrotik -from homeassistant.components.mikrotik.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -84,4 +83,3 @@ async def test_unload_entry(hass: HomeAssistant) -> None: await hass.async_block_till_done() assert entry.state is ConfigEntryState.NOT_LOADED - assert entry.entry_id not in hass.data[DOMAIN]